[TRTLLM-12671][chore] Refactor visual_gen's attention handling; Fix LTX-2 & Cosmos3

tensorrt_llm/_torch/attention_backend/fmha/fallback.py

@@ -109,7 +109,8 @@ def forward(
             max_context_length=metadata.max_context_length,
             max_seq_len=metadata.max_seq_len,
             trtllm_gen_jit_warmup=metadata.trtllm_gen_jit_warmup,
-            is_cross=metadata.is_cross,
+            # Selects the KV layout, not merely unequal Q/KV lengths.
+            is_cross=metadata.is_cross_with_kv_cache,
             # --- Per-call (AttentionForwardArgs) ---
             out_scale=forward_args.out_scale,
             kv_scale_orig_quant=forward_args.kv_scale_orig_quant,

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -86,6 +86,16 @@ class TrtllmAttentionMetadata(AttentionMetadata):
     # when beam search is enabled.
     beam_width: int = 1
 
+    @property
+    def is_cross_with_kv_cache(self) -> bool:
+        """Cross-attention whose K/V is read from a KV-cache pool.
+
+        thop uses ``is_cross`` flag to handle cross attention with kv cache
+        (``Q_CONTIGUOUS_KV`` / ``Q_PAGED_KV``, fed by ``cross_kv``), while
+        cross-attention with ``SEPARATE_QKV`` is handled without the flag.
+        """
+        return self.is_cross and self.kv_cache_manager is not None
+
     @property
     def effective_beam_width(self) -> int:
         # Only use this for the fallback kernel's beam_width argument.
@@ -1798,7 +1808,13 @@ def forward(
                 and metadata._seq_lens_cuda is not None):
             metadata.update_blackwell_first_sparse_mask_offset()
 
-        if metadata.is_cross:
+        # The ``cacheless`` (separate-QKV) cross attention does not need cross-
+        # attention specific handling
+        _cacheless_cross = (metadata.is_cross
+                            and metadata.kv_cache_manager is None
+                            and k is not None and v is not None)
+
+        if metadata.is_cross and not _cacheless_cross:
             if k is not None and v is not None:
                 k_flat = k.contiguous().view(k.shape[0], -1)
                 v_flat = v.contiguous().view(v.shape[0], -1)

tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md

@@ -29,6 +29,11 @@ end-to-end workflow.
    T2V / TI2V) with LPIPS — or an equivalent metric — at an explicit
    threshold, comparing against a baseline image/video produced by a
    reference framework.
+6. **Attention metadata is a forward argument.** Models take metadata as
+   required `forward()` parameters; the pipeline builds it per forward
+   from the model's `attn_metadata_spec()`. Each attention *site*
+   (e.g., LTX-2's Vid2Vid, Vid2Sound, etc.), with its own Q/KV lengths,
+   gets its own metadata object.
 
 ## 2. API
 

tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py

@@ -28,7 +28,7 @@
 from tensorrt_llm.logger import logger
 from tensorrt_llm.visual_gen.args import QuantAttentionConfig
 
-from ....attention_backend.interface import PredefinedAttentionMask
+from ....attention_backend.interface import AttentionMetadata, PredefinedAttentionMask
 from ..interface import AttentionBackend, AttentionTensorLayout
 
 _cute_dsl_import_error: BaseException | None = None
@@ -740,6 +740,7 @@ def forward(
         k: torch.Tensor,
         v: torch.Tensor,
         *,
+        attn_metadata: AttentionMetadata,
         attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL,
         **kwargs,
     ) -> torch.Tensor:
@@ -752,19 +753,26 @@ def forward(
             q: Query tensor [batch_size, seq_len, num_heads, head_dim]
             k: Key tensor [batch_size, seq_len_kv, num_kv_heads, head_dim]
             v: Value tensor [batch_size, seq_len_kv, num_kv_heads, head_dim]
+            attn_metadata: Unused; the CuTe DSL kernels derive everything from
+                tensor shapes. Preserved for future metadata-consuming variants
+                of this path.
             attention_mask: Attention mask type (CAUSAL or FULL)
 
         Returns:
             Output tensor [batch_size, seq_len, num_heads, head_dim]
         """
-        output, _ = self.forward_with_lse(q, k, v, attention_mask=attention_mask, **kwargs)
+        output, _ = self.forward_with_lse(
+            q, k, v, attn_metadata=attn_metadata, attention_mask=attention_mask, **kwargs
+        )
         return output
 
     def forward_with_lse(
         self,
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
+        *,
+        attn_metadata: AttentionMetadata,
         attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL,
         **kwargs,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -777,12 +785,20 @@ def forward_with_lse(
                     always in float32. Used for numerically stable combination of
                     partial attention results in Attention2D parallelism.
         """
+        _ = attn_metadata
         q, k, v, is_causal, origin_dtype = self._prepare_inputs(q, k, v, attention_mask)
         output, lse = self._fwd(q, k, v, is_causal, **kwargs)
         if output.dtype != origin_dtype:
             output = output.to(origin_dtype)
         return output, lse.transpose(1, 2)
 
+    @property
+    def requires_metadata(self) -> bool:
+        """While CuTe's kernel supports varlen, its not wired to CuTeDSLAttention.
+
+        No metadata referred."""
+        return False
+
     @classmethod
     def support_lse(cls) -> bool:
         return True

tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py

@@ -30,6 +30,7 @@
 import torch
 import torch.nn.functional as F
 
+from ....attention_backend.interface import AttentionMetadata
 from ..interface import AttentionBackend, AttentionTensorLayout
 
 _vsa_import_error = None
@@ -300,13 +301,18 @@ def forward(
         k: torch.Tensor,
         v: torch.Tensor,
         *,
+        attn_metadata: AttentionMetadata,
         gate_compress: Optional[torch.Tensor] = None,
         gate_fine: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
         """
         VSA forward: coarse mean-pool + fine block-sparse top-K.
 
+        VSA carries its own sparse-layout metadata through
+        ``set_vsa_forward_context``; ``attn_metadata`` is unused here and
+        preserved for future metadata-consuming variants of this path.
+
         Args:
             q, k, v: [B, S, H, D] in original (un-tiled) token order.
             gate_compress: [B, S, H, D] G_c gate weighting the coarse branch O_c.
@@ -316,6 +322,7 @@ def forward(
         Returns:
             [B, S, H, D] in the same original token order.
         """
+        _, _ = attn_metadata, kwargs
         if gate_compress is None:
             raise ValueError(
                 "VSAAttention requires gate_compress. "
@@ -399,6 +406,11 @@ def forward(
             return gate_compress * o_c_full + gate_fine * o_f
         return gate_compress * o_c_full + o_f
 
+    @property
+    def requires_metadata(self) -> bool:
+        """Currently VSAAttention only supports non-varlen batching. No metadata referred."""
+        return False
+
     @classmethod
     def support_lse(cls) -> bool:
         return False

tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py

@@ -24,7 +24,7 @@
 
 import torch
 
-from ...attention_backend.interface import PredefinedAttentionMask
+from ...attention_backend.interface import AttentionMetadata, PredefinedAttentionMask
 from .interface import AttentionBackend, AttentionTensorLayout
 
 _flash_attn_fwd_import_error = None
@@ -122,6 +122,7 @@ def forward(
         k: torch.Tensor,
         v: torch.Tensor,
         *,
+        attn_metadata: AttentionMetadata,
         attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL,
         key_padding_mask: Optional[torch.Tensor] = None,
         **kwargs,
@@ -135,6 +136,7 @@ def forward(
             q: Query tensor [batch_size, seq_len, num_heads, head_dim]
             k: Key tensor [batch_size, seq_len_kv, num_kv_heads, head_dim]
             v: Value tensor [batch_size, seq_len_kv, num_kv_heads, head_dim]
+            attn_metadata: Unused; FA4 derives everything from tensor shapes.
             attention_mask: Attention mask type (CAUSAL or FULL)
             key_padding_mask: Optional ``[B, S_kv]`` bool tensor; True = valid,
                 False = pad. Translated to FA4's ``seqused_k = mask.sum(dim=1)``
@@ -147,6 +149,7 @@ def forward(
             q,
             k,
             v,
+            attn_metadata=attn_metadata,
             attention_mask=attention_mask,
             key_padding_mask=key_padding_mask,
             **kwargs,
@@ -158,6 +161,8 @@ def forward_with_lse(
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
+        *,
+        attn_metadata: AttentionMetadata,
         attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL,
         key_padding_mask: Optional[torch.Tensor] = None,
         **kwargs,
@@ -171,6 +176,7 @@ def forward_with_lse(
                     always in float32. Used for numerically stable combination of
                     partial attention results in Attention2D parallelism.
         """
+        _, _ = attn_metadata, kwargs
         q, k, v, is_causal, origin_dtype = self._prepare_inputs(q, k, v, attention_mask)
         seqused_k = None
         if key_padding_mask is not None:
@@ -190,6 +196,11 @@ def forward_with_lse(
             output = output.to(origin_dtype)
         return output, lse
 
+    @property
+    def requires_metadata(self) -> bool:
+        """Currently FlashAttn4Attention only supports non-varlen batching. No metadata referred."""
+        return False
+
     @classmethod
     def support_lse(cls) -> bool:
         return True

tensorrt_llm/_torch/visual_gen/attention_backend/interface.py

@@ -20,9 +20,12 @@
 
 from abc import ABC, abstractmethod
 from enum import Enum
+from typing import Type
 
 import torch
 
+from ...attention_backend.interface import AttentionMetadata
+
 
 class AttentionTensorLayout(str, Enum):
     """
@@ -42,8 +45,17 @@ class AttentionBackend(ABC):
     Every backend must implement ``forward`` and declare a ``preferred_layout``.
     Backends pick the kwargs they need from the caller and ignore the rest
     via ``**kwargs``.
+
+    ``attn_metadata`` is a required argument on every forward, mirroring the LLM
+    backend. Backends that do not consume metadata (VANILLA, FA4, CuTe DSL) accept
+    and ignore it; only the TRTLLM backend reads it. See
+    ``visual_gen/attention_backend/metadata.py`` for the site contract.
     """
 
+    #: Concrete metadata type this backend consumes; callers construct from
+    #: it, as the LLM engine does.
+    Metadata: Type[AttentionMetadata] = AttentionMetadata
+
     def __call__(self, *args, **kwargs) -> torch.Tensor:
         return self.forward(*args, **kwargs)
 
@@ -53,18 +65,30 @@ def forward(
         q: torch.Tensor,
         k: torch.Tensor | None = None,
         v: torch.Tensor | None = None,
+        *,
+        attn_metadata: AttentionMetadata,
         **kwargs,
     ) -> torch.Tensor: ...
 
     @property
     @abstractmethod
     def preferred_layout(self) -> AttentionTensorLayout: ...
 
+    @property
+    def requires_metadata(self) -> bool:
+        """Queries if the configured attention module requires a prepared metadata to run.
+
+        For size-aligned diffusion models, many attention backends can run without a metadata.
+        """
+        return True
+
     def forward_with_lse(
         self,
         q: torch.Tensor,
         k: torch.Tensor | None = None,
         v: torch.Tensor | None = None,
+        *,
+        attn_metadata: AttentionMetadata,
         **kwargs,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         raise NotImplementedError(