[None][fix] LTX-2: re-enable Ulysses for v2a cross-attention (#15303)

Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

Author: luyiyun1021

tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py

@@ -812,6 +812,7 @@ def wrap_parallel_attention(
     *,
     visual_gen_mapping: Optional["VisualGenMapping"] = None,
     enable_sequence_parallel: bool = True,
+    use_ulysses: bool = True,
     async_ulysses: bool = False,
 ) -> AttentionBackend:
     """Wrap a compute backend with the configured parallelism strategy.
@@ -822,6 +823,12 @@ def wrap_parallel_attention(
 
     When ``enable_sequence_parallel`` is False, no wrappers are applied (callers
     use this for cross-attention paths that cannot use Ulysses/Ring/Attention2D).
+
+    ``use_ulysses`` gates the Ulysses head-sharding wrap independently of
+    ``ulysses_size``: pass False to skip it even when ``ulysses_size > 1`` (e.g. a
+    SEPARATE_QKV cross-attn that falls back to all-gather and built its inner
+    backend with the full, un-sharded head count). This keeps the wrap consistent
+    with the caller's inner head-count decision.
     """
     if not enable_sequence_parallel or visual_gen_mapping is None:
         return attn
@@ -840,7 +847,7 @@ def wrap_parallel_attention(
     elif ring_size > 1:
         attn = RingAttention(attn, process_group=vgm.ring_group)
 
-    if ulysses_size > 1:
+    if ulysses_size > 1 and use_ulysses:
         attn = UlyssesAttention(
             attn,
             process_group=vgm.ulysses_group,

tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py

@@ -170,11 +170,14 @@ def __init__(
             # (sharded inner backend + UlyssesAttention) for both self-attn and
             # cross-attn paths.
 
-        # For audio self-attention that may need a runtime Ulysses toggle
-        # (sequence length not always divisible by ulysses_size), create a
-        # plain backend as fallback.  The base class already set self.attn
-        # to UlyssesAttention(inner_backend=sharded_backend).
-        if use_ulysses and not self._is_cross_attn and ulysses_size > 1:
+        # Build a Ulysses/plain dual-attn pair so set_ulysses_active() can toggle
+        # at runtime. Needed for self-attn (audio seq not always divisible by
+        # ulysses_size) and for v2a cross-attn under pure Ulysses (cp_size == 1),
+        # where it lets the block forward use Ulysses a2a instead of all-gathering
+        # the full video K/V. Combined ring/attn2d + Ulysses (cp_size > 1) keeps
+        # cross-attn on the all-gather fallback. The base class already set
+        # self.attn to UlyssesAttention(inner_backend=sharded_backend).
+        if use_ulysses and (cp_size == 1 or not self._is_cross_attn) and ulysses_size > 1:
             self._ulysses_attn = self.attn
             self._plain_attn = create_attention(
                 backend=self.attn_backend,
@@ -684,6 +687,7 @@ def _init_av_cross_modules(self, v_cfg, a_cfg, rope_type, eps, model_config, idx
             layer_idx=idx,
             module_name=f"transformer_blocks.{idx}.video_to_audio_attn",
             enable_sequence_parallel=True,
+            use_ulysses=True,
         )
         self.scale_shift_table_a2v_ca_audio = nn.Parameter(torch.empty(5, a_cfg.dim))
         self.scale_shift_table_a2v_ca_video = nn.Parameter(torch.empty(5, v_cfg.dim))

tensorrt_llm/_torch/visual_gen/modules/attention.py

@@ -4,6 +4,7 @@
 import torch
 import torch.nn as nn
 
+from tensorrt_llm.logger import logger
 from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig
 
 from ...modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
@@ -182,11 +183,22 @@ def __init__(
         # The async-ulysses path uses SEPARATE_QKV for stream-pipelined
         # V/Q/K projections AND still needs the head-sharding wrap — opt in
         # via async_ulysses=True.
+        cp_size = vgm.cp_size if vgm else 1
         use_ulysses = (
             ulysses_size > 1
             and enable_sequence_parallel
-            and (self.qkv_mode != QKVMode.SEPARATE_QKV or async_ulysses)
+            and (self.qkv_mode != QKVMode.SEPARATE_QKV or async_ulysses or cp_size == 1)
         )
+        if ulysses_size > 1 and enable_sequence_parallel and not use_ulysses:
+            # Ulysses was requested (ulysses_size > 1, SP on) but disabled: this is a
+            # SEPARATE_QKV cross-attention that is neither async nor pure-Ulysses
+            # (cp_size > 1), so it falls back to the all-gather K/V path.
+            logger.debug(
+                f"Attention(layer={layer_idx}): Ulysses disabled despite ulysses_size="
+                f"{ulysses_size} — qkv_mode={self.qkv_mode.value}, "
+                f"async_ulysses={async_ulysses}, cp_size={cp_size} "
+                f"(SEPARATE_QKV cross-attn needs async_ulysses or cp_size==1)."
+            )
 
         # Compute head counts for the backend
         # Ulysses shards heads across workers; inner backend sees sharded count
@@ -234,6 +246,7 @@ def __init__(
             self.attn,
             visual_gen_mapping=vgm,
             enable_sequence_parallel=enable_sequence_parallel,
+            use_ulysses=use_ulysses,
             async_ulysses=use_ulysses and async_ulysses,
         )