[https://nvbugs/6705472][feat] FA4 variable-length cross-attention

tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py

@@ -141,13 +141,25 @@ def _fwd(
         v: torch.Tensor,
         causal: bool,
         seqused_k: Optional[torch.Tensor] = None,
+        cu_seqlens_q: Optional[torch.Tensor] = None,
+        cu_seqlens_k: Optional[torch.Tensor] = None,
+        max_seqlen_q: Optional[int] = None,
+        max_seqlen_k: Optional[int] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Calls _flash_attn_fwd with torch.compile disabled. Returns (output, lse)."""
+        """Calls _flash_attn_fwd with torch.compile disabled. Returns (output, lse).
+
+        cu_seqlens_q/cu_seqlens_k switch the kernel into ragged mode: q/k/v are
+        then expected pre-packed as [total_tokens, H, D] instead of [B, S, H, D].
+        """
         # FA4's private forward API may append diagnostics that this backend does not consume.
         output, lse, *_ = _flash_attn_fwd(
             q,
             k,
             v,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
             seqused_k=seqused_k,
             softmax_scale=self.scale,
             causal=causal,
@@ -230,18 +242,47 @@ def forward_with_lse(
         v: torch.Tensor,
         attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL,
         key_padding_mask: Optional[torch.Tensor] = None,
+        cu_seqlens_q: Optional[torch.Tensor] = None,
+        cu_seqlens_kv: Optional[torch.Tensor] = None,
+        max_seqlen_q: Optional[int] = None,
+        max_seqlen_kv: Optional[int] = None,
         **kwargs,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         Forward pass returning both output and log-sum-exp (LSE).
 
         Returns:
-            output: [batch_size, seq_len, num_heads, head_dim]
-            lse:    [batch_size, num_heads, seq_len] — log-sum-exp per query position,
-                    always in float32. Used for numerically stable combination of
-                    partial attention results in Attention2D parallelism.
+            output: [batch_size, seq_len, num_heads, head_dim] if cu_seqlens_q is
+                unset, else [total_q_tokens, num_heads, head_dim].
+            lse:    [batch_size, num_heads, seq_len] if cu_seqlens_q is unset, else
+                [num_heads, total_q_tokens]. Callers combining LSE across ranks
+                need to handle both shapes.
         """
         q, k, v, is_causal, origin_dtype = self._prepare_inputs(q, k, v, attention_mask)
+
+        if cu_seqlens_kv is not None:
+            assert key_padding_mask is None, (
+                "cu_seqlens_kv (ragged varlen) and key_padding_mask (padded+mask) "
+                "are mutually exclusive attention modes"
+            )
+            assert max_seqlen_kv is not None, "cu_seqlens_kv requires max_seqlen_kv"
+            assert cu_seqlens_q is None or max_seqlen_q is not None, (
+                "cu_seqlens_q requires max_seqlen_q"
+            )
+            output, lse = self._fwd(
+                q,
+                k,
+                v,
+                is_causal,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_kv,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_kv,
+            )
+            if output.dtype != origin_dtype:
+                output = output.to(origin_dtype)
+            return output, lse
+
         seqused_k = None
         if key_padding_mask is not None:
             assert not is_causal, "key_padding_mask is not supported with causal attention"
@@ -260,6 +301,10 @@ def forward_with_lse(
             output = output.to(origin_dtype)
         return output, lse
 
+    @classmethod
+    def supports_varlen(cls) -> bool:
+        return True
+
     @classmethod
     def support_lse(cls) -> bool:
         return True

tensorrt_llm/_torch/visual_gen/attention_backend/interface.py

@@ -80,3 +80,8 @@ def support_fused_qkv(cls) -> bool:
     def support_lse(cls) -> bool:
         """Whether the backend supports returning the softmax log-sum-exp (LSE) of the attention weights."""
         return False
+
+    @classmethod
+    def supports_varlen(cls) -> bool:
+        """Whether the backend accepts ragged K/V via ``cu_seqlens_kv``/``max_seqlen_kv``."""
+        return False

tensorrt_llm/_torch/visual_gen/config.py

@@ -747,7 +747,7 @@ def from_pretrained(
 
         attention_metadata_state = (
             create_attention_metadata_state()
-            if attention_cfg.backend in ("TRTLLM", "FLASHINFER")
+            if attention_cfg.backend in ("TRTLLM", "FLASHINFER", "FA4")
             else None
         )
 

tensorrt_llm/_torch/visual_gen/modules/attention.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -151,6 +151,7 @@ def __init__(
         )
 
         attention_metadata_state = getattr(config, "attention_metadata_state", None)
+        self._metadata_state = attention_metadata_state
 
         if self.qk_norm:
             # "full": norm over all heads combined (e.g. WAN, dim=q_dim)
@@ -269,6 +270,9 @@ def __init__(
             async_ulysses=use_ulysses and async_ulysses,
         )
 
+        # Checked post-wrap so this reflects self.attn as it's actually called.
+        self.supports_varlen = self.attn.supports_varlen()
+
     @staticmethod
     def _qualified_module_name(
         component_name: Optional[str],
@@ -541,7 +545,13 @@ def _attn_impl(
         Two layout paths:
         1. HND backends (VANILLA): [B, S, H*D] -> [B, H, S, D]
         2. NHD backends (TRTLLM, UlyssesAttention, Attention2DAttention): [B, S, H*D] -> [B, S, H, D]
+
+        A third path (see ``_attn_impl_varlen_kv``) handles ragged K/V when the
+        caller passes ``cu_seqlens_kv``.
         """
+        if kwargs.get("cu_seqlens_kv") is not None:
+            return self._attn_impl_varlen_kv(q, k, v, **kwargs)
+
         backend_layout = getattr(self.attn, "preferred_layout", AttentionTensorLayout.NHD)
 
         batch_size = q.shape[0]
@@ -589,6 +599,91 @@ def _reshape_gate(gate: torch.Tensor) -> torch.Tensor:
         else:
             return out.flatten(2)
 
+    @staticmethod
+    def _cu_seqlens_kv_and_max(
+        metadata_state: Optional[Dict[str, Any]],
+        kv_lens: Tuple[int, ...],
+        device: torch.device,
+    ) -> Tuple[torch.Tensor, int]:
+        """Depends only on kv_lens (constant across layers/denoising steps),
+        not K/V content, so it's cached in the model-scoped metadata_state dict."""
+        cache_key = (kv_lens, device)
+        if metadata_state is not None:
+            cache = metadata_state.setdefault("varlen_kv_cache", {})
+            cached = cache.get(cache_key)
+            if cached is not None:
+                return cached
+
+        cu_list = [0]
+        for n in kv_lens:
+            cu_list.append(cu_list[-1] + n)
+        cu_seqlens_kv = torch.tensor(cu_list, dtype=torch.int32, device=device)
+        max_seqlen_kv = max(kv_lens)
+        result = (cu_seqlens_kv, max_seqlen_kv)
+
+        if metadata_state is not None:
+            cache[cache_key] = result
+        return result
+
+    @staticmethod
+    def pack_ragged_kv(
+        k: torch.Tensor,
+        v: torch.Tensor,
+        kv_lens: List[int],
+        metadata_state: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+        """Slice each sample's true K/V rows out of padded [B, S, H*D] K/V
+        into [total_kv_tokens, H*D], plus cu_seqlens_kv/max_seqlen_kv for
+        _attn_impl_varlen_kv. kv_lens is a plain host-side list; a CUDA
+        tensor would force a sync on .tolist().
+        """
+        k_parts, v_parts = [], []
+        for i, n in enumerate(kv_lens):
+            k_parts.append(k[i, :n])
+            v_parts.append(v[i, :n])
+        k_ragged = torch.cat(k_parts, dim=0)
+        v_ragged = torch.cat(v_parts, dim=0)
+        cu_seqlens_kv, max_seqlen_kv = Attention._cu_seqlens_kv_and_max(
+            metadata_state, tuple(kv_lens), k.device
+        )
+        return k_ragged, v_ragged, cu_seqlens_kv, max_seqlen_kv
+
+    def _attn_impl_varlen_kv(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Ragged K/V cross-attention: Q stays batched [B, S, H*D], K/V arrive
+        pre-packed as [total_kv_tokens, H_kv*D]."""
+        if not self.supports_varlen:
+            raise ValueError(
+                f"{type(self.attn).__name__} does not support varlen cross-attention "
+                "(cu_seqlens_kv); check `Attention.supports_varlen` before passing "
+                "ragged K/V, and fall back to the padded K/V path otherwise."
+            )
+
+        cu_seqlens_kv = kwargs.pop("cu_seqlens_kv")
+        max_seqlen_kv = kwargs.pop("max_seqlen_kv")
+
+        batch_size, seq_len_q = q.shape[0], q.shape[1]
+
+        q = q.reshape(batch_size, seq_len_q, self.local_num_attention_heads, self.head_dim)
+        k = k.reshape(-1, self.local_num_key_value_heads, self.head_dim)
+        v = v.reshape(-1, self.local_num_key_value_heads, self.head_dim)
+
+        kwargs.update(
+            {
+                "batch_size": batch_size,
+                "seq_len": seq_len_q,
+                "cu_seqlens_kv": cu_seqlens_kv,
+                "max_seqlen_kv": max_seqlen_kv,
+            }
+        )
+        out = self.attn.forward(q=q, k=k, v=v, **kwargs)
+        return out.reshape(batch_size, seq_len_q, -1)
+
     def forward(
         self,
         hidden_states: torch.Tensor | Fp4QuantizedTensor,

tests/integration/test_lists/test-db/l0_b200.yml

@@ -281,6 +281,7 @@ l0_b200:
   - unittest/_torch/visual_gen/test_attention_cudnn.py
   - unittest/_torch/visual_gen/test_attention_integration.py
   - unittest/_torch/visual_gen/test_attention_fa4.py
+  - unittest/_torch/visual_gen/test_varlen_attention.py
   - unittest/_torch/visual_gen/test_attention_perf.py
   - unittest/_torch/visual_gen/test_qwen_image_layered_registry.py
   - unittest/_torch/visual_gen/test_trtllm_serve_e2e.py