From 171d8d9418429c5911afc7f40764b98f2b0ef3f9 Mon Sep 17 00:00:00 2001
From: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Date: Tue, 8 Sep 2026 11:49:47 -0700
Subject: [PATCH 1/9] add FA4 cu_seqlens varlen cross-attention support

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
---
 .../attention_backend/flash_attn4.py          |  55 ++-
 .../visual_gen/attention_backend/interface.py |   5 +
 .../_torch/visual_gen/modules/attention.py    |  70 ++++
 tensorrt_llm/visual_gen/args.py               |   8 +
 .../visual_gen/test_varlen_attention.py       | 341 ++++++++++++++++++
 .../_torch/visual_gen/test_wan_transformer.py |  96 ++++-
 6 files changed, 569 insertions(+), 6 deletions(-)
 create mode 100644 tests/unittest/_torch/visual_gen/test_varlen_attention.py

diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py
index b65eb1e9a645..f9e033a2978e 100644
--- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py
+++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py
@@ -72,13 +72,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,
@@ -161,18 +173,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], or
+                [total_q_tokens, num_heads, head_dim] in the varlen path
+                (cu_seqlens_kv set).
+            lse:    [batch_size, num_heads, seq_len] in the padded path, or
+                [num_heads, total_q_tokens] in the varlen path. Callers doing
+                Attention2D/Ring LSE-based combination 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 (
+                cu_seqlens_q is not None and max_seqlen_q is not None and max_seqlen_kv is not None
+            ), "cu_seqlens_kv requires cu_seqlens_q, max_seqlen_q, and max_seqlen_kv"
+            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"
@@ -191,6 +232,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
diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py b/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py
index 7a866acf5879..6837e00bb262 100644
--- a/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py
+++ b/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
diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py
index 8c19484229c7..1d7552339c86 100644
--- a/tensorrt_llm/_torch/visual_gen/modules/attention.py
+++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py
@@ -3,6 +3,7 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 from tensorrt_llm.logger import logger
 from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig
@@ -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,66 @@ def _reshape_gate(gate: torch.Tensor) -> torch.Tensor:
         else:
             return out.flatten(2)
 
+    @staticmethod
+    def pack_ragged_kv(
+        k: torch.Tensor, v: torch.Tensor, kv_lens: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """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 for _attn_impl_varlen_kv.
+        """
+        k_parts, v_parts = [], []
+        for i, n in enumerate(kv_lens.tolist()):
+            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 = F.pad(torch.cumsum(kv_lens, dim=0), (1, 0)).to(torch.int32)
+        return k_ragged, v_ragged, cu_seqlens_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 is uniform-length [B, S, H*D], K/V arrive
+        pre-packed as [total_kv_tokens, H_kv*D]. The caller builds
+        ``cu_seqlens_kv``/``max_seqlen_kv`` from the real per-sample lengths.
+        """
+        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]
+        total_q = batch_size * seq_len_q
+        cu_seqlens_q = torch.arange(
+            0, total_q + seq_len_q, seq_len_q, dtype=torch.int32, device=q.device
+        )
+
+        q = q.reshape(total_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_q": cu_seqlens_q,
+                "cu_seqlens_kv": cu_seqlens_kv,
+                "max_seqlen_q": seq_len_q,
+                "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,
diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py
index f7d483389dd0..f65c6a39135e 100644
--- a/tensorrt_llm/visual_gen/args.py
+++ b/tensorrt_llm/visual_gen/args.py
@@ -119,6 +119,14 @@ class AttentionConfig(StrictBaseModel):
             "skip_softmax (TRTLLM / CUTEDSL backends) or VSA (CUTEDSL backend)."
         ),
     )
+    enable_varlen_cfg: bool = Field(
+        False,
+        status="prototype",
+        description=(
+            "Pack unequal-length CFG text cross-attention via cu_seqlens instead of "
+            "padding. Requires an FA4 backend. Not yet wired through any model."
+        ),
+    )
 
     @model_validator(mode="after")
     def _validate_quant_attention_config(self) -> "AttentionConfig":
diff --git a/tests/unittest/_torch/visual_gen/test_varlen_attention.py b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
new file mode 100644
index 000000000000..fc01b8752283
--- /dev/null
+++ b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
@@ -0,0 +1,341 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical parity for the FA4 varlen (cu_seqlens) cross-attention capability.
+
+Covers two layers:
+  1. ``FlashAttn4Attention.forward_with_lse`` varlen path -- the Blackwell-tuned
+     ``flash_attn.cute`` kernel, packed ragged K/V vs. a per-sample SDPA
+     reference. Requires CUDA + the FA4 CuTe kernel.
+  2. ``Attention._attn_impl_varlen_kv`` -- the dispatch/reshape glue in
+     ``modules/attention.py`` that builds ``cu_seqlens_q`` and reshapes Q/K/V
+     around the backend call. Requires CUDA (delegates to layer 1).
+
+This capability is not wired into any model yet; these tests exercise the
+backend and dispatch layer directly.
+"""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.attention_backend.interface import PredefinedAttentionMask
+from tensorrt_llm._torch.visual_gen.attention_backend import (
+    CuTeDSLAttention,
+    TrtllmAttention,
+    VanillaAttention,
+)
+from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import (
+    FlashAttn4Attention,
+    _flash_attn_fwd,
+)
+from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig
+from tensorrt_llm._torch.visual_gen.modules import attention as attention_module
+from tensorrt_llm._torch.visual_gen.modules.attention import Attention
+from tensorrt_llm.visual_gen.args import AttentionConfig
+
+FA4_AVAILABLE = _flash_attn_fwd is not None
+
+fa4_cuda_only = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="FA4 requires CUDA"),
+    pytest.mark.skipif(not FA4_AVAILABLE, reason="FA4 kernel not available"),
+]
+
+
+def _sdpa_reference(q, k_list, v_list, scale):
+    """Per-sample SDPA over each sample's true (un-padded) K/V, looped -- the
+    same reference structure the LLM-side vanilla.py fallback uses. q is
+    uniform-length [B, H, S_q, D]; k_list/v_list are per-sample [H_kv, S_kv_i, D].
+    """
+    outs = []
+    for i in range(q.shape[0]):
+        out_i = torch.nn.functional.scaled_dot_product_attention(
+            q[i : i + 1], k_list[i].unsqueeze(0), v_list[i].unsqueeze(0), scale=scale
+        )
+        outs.append(out_i)
+    return torch.cat(outs, dim=0)
+
+
+def _run_packed(attn, q_bhsd, k_list, v_list, S_q, H, d_h, device):
+    """Pack q_bhsd/k_list/v_list (whatever subset is passed in) and run the
+    varlen path. Returns [B, S_q, H, d_h]."""
+    B = q_bhsd.shape[0]
+    kv_lens = [k.shape[1] for k in k_list]
+    q_packed = q_bhsd.transpose(1, 2).reshape(B * S_q, H, d_h)
+    cu_seqlens_q = torch.arange(0, (B + 1) * S_q, S_q, dtype=torch.int32, device=device)
+    k_packed = torch.cat([k.transpose(0, 1) for k in k_list], dim=0)
+    v_packed = torch.cat([v.transpose(0, 1) for v in v_list], dim=0)
+    lens_tensor = torch.tensor(kv_lens, dtype=torch.int32, device=device)
+    cu_seqlens_kv = torch.nn.functional.pad(torch.cumsum(lens_tensor, dim=0), (1, 0)).to(
+        torch.int32
+    )
+    out = attn.forward(
+        q=q_packed,
+        k=k_packed,
+        v=v_packed,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_kv=cu_seqlens_kv,
+        max_seqlen_q=S_q,
+        max_seqlen_kv=max(kv_lens),
+    )
+    return out.reshape(B, S_q, H, d_h)
+
+
+def _run_ragged_kv_vs_sdpa(attn, B, S_q, H, d_h, kv_lens, device, dtype):
+    """Shared harness: pack ragged K/V, run attn's varlen path, compare
+    against the per-sample SDPA reference. Returns (out_reshaped, ref)."""
+    torch.manual_seed(1)
+    q_bhsd = torch.randn(B, H, S_q, d_h, device=device, dtype=dtype)
+    k_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+    v_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+
+    ref = _sdpa_reference(q_bhsd, k_list, v_list, attn.scale)  # [B, H, S_q, d_h]
+
+    out = _run_packed(attn, q_bhsd, k_list, v_list, S_q, H, d_h, device)
+    out_reshaped = out.transpose(1, 2)  # [B, H, S_q, D]
+    return out_reshaped, ref
+
+
+@pytest.mark.parametrize("backend_cls", [VanillaAttention, TrtllmAttention, CuTeDSLAttention])
+def test_backend_without_varlen_support_defaults_false(backend_cls):
+    assert backend_cls.supports_varlen() is False
+
+
+def test_supports_varlen_checked_post_wrap(monkeypatch):
+    """Must reflect self.attn after wrap_parallel_attention, not before."""
+    fake_wrapped = type("FakeWrapped", (), {"supports_varlen": staticmethod(lambda: False)})()
+    monkeypatch.setattr(
+        attention_module, "wrap_parallel_attention", lambda backend, **kw: fake_wrapped
+    )
+
+    attn = Attention(
+        hidden_size=64,
+        num_attention_heads=4,
+        head_dim=16,
+        config=DiffusionModelConfig(attention=AttentionConfig(backend="FA4")),
+    )
+    assert attn.supports_varlen is False
+
+
+class TestFA4VarlenKv:
+    """FlashAttn4Attention varlen path (Blackwell-tuned flash_attn.cute kernel)
+    vs. per-sample SDPA reference."""
+
+    pytestmark = fa4_cuda_only
+
+    def test_ragged_kv_matches_padded_reference(self):
+        device, dtype = "cuda", torch.bfloat16
+        B, S_q, H, d_h = 2, 4, 8, 64
+        kv_lens = [3, 9]
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+        out, ref = _run_ragged_kv_vs_sdpa(attn, B, S_q, H, d_h, kv_lens, device, dtype)
+        torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2)
+
+    def test_supports_varlen(self):
+        assert FlashAttn4Attention.supports_varlen() is True
+
+    def test_split_consistency(self):
+        """Running a sub-batch on its own must match its slice of a larger
+        packed batch -- catches batch-offset/indexing bugs in cu_seqlens
+        dispatch that a single-batch SDPA comparison could miss."""
+        device, dtype = "cuda", torch.bfloat16
+        B, S_q, H, d_h = 3, 4, 8, 64
+        kv_lens = [3, 9, 5]
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+
+        torch.manual_seed(3)
+        q_bhsd = torch.randn(B, H, S_q, d_h, device=device, dtype=dtype)
+        k_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        v_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+
+        full_out = _run_packed(attn, q_bhsd, k_list, v_list, S_q, H, d_h, device)
+
+        sub_idx = [1, 2]
+        sub_out = _run_packed(
+            attn,
+            q_bhsd[sub_idx],
+            [k_list[i] for i in sub_idx],
+            [v_list[i] for i in sub_idx],
+            S_q,
+            H,
+            d_h,
+            device,
+        )
+
+        torch.testing.assert_close(sub_out, full_out[sub_idx], rtol=2e-2, atol=2e-2)
+
+    def test_uneven_boundary_lengths(self):
+        """One sample at max_seqlen_kv (no padding to remove), the other short."""
+        device, dtype = "cuda", torch.bfloat16
+        B, S_q, H, d_h = 2, 4, 8, 64
+        kv_lens = [2, 64]
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+        out, ref = _run_ragged_kv_vs_sdpa(attn, B, S_q, H, d_h, kv_lens, device, dtype)
+        torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2)
+
+    def test_zero_length_kv_sample(self):
+        """Empty CFG branch (e.g. empty negative prompt) -- pins actual FA4 behavior."""
+        device, dtype = "cuda", torch.bfloat16
+        B, S_q, H, d_h = 2, 4, 8, 64
+        kv_lens = [0, 9]
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+        torch.manual_seed(1)
+        q_bhsd = torch.randn(B, H, S_q, d_h, device=device, dtype=dtype)
+        k_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        v_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        out = _run_packed(attn, q_bhsd, k_list, v_list, S_q, H, d_h, device)
+        assert not torch.isnan(out).any()
+
+    def test_all_equal_lengths(self):
+        """Degenerate case: nothing to pack, every sample the same length."""
+        device, dtype = "cuda", torch.bfloat16
+        B, S_q, H, d_h = 3, 4, 8, 64
+        kv_lens = [16, 16, 16]
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+        out, ref = _run_ragged_kv_vs_sdpa(attn, B, S_q, H, d_h, kv_lens, device, dtype)
+        torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2)
+
+    def test_raises_when_causal_combined_with_key_padding_mask(self):
+        device, dtype = "cuda", torch.bfloat16
+        H, d_h = 4, 32
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+        q = torch.randn(2, 3, H, d_h, device=device, dtype=dtype)
+        k = torch.randn(2, 3, H, d_h, device=device, dtype=dtype)
+        v = torch.randn(2, 3, H, d_h, device=device, dtype=dtype)
+        with pytest.raises(AssertionError, match="key_padding_mask is not supported"):
+            attn.forward_with_lse(
+                q,
+                k,
+                v,
+                attention_mask=PredefinedAttentionMask.CAUSAL,
+                key_padding_mask=torch.ones(2, 3, dtype=torch.bool, device=device),
+            )
+
+    def test_raises_when_key_padding_mask_combined_with_cu_seqlens(self):
+        device, dtype = "cuda", torch.bfloat16
+        H, d_h = 4, 32
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+        q = torch.randn(6, H, d_h, device=device, dtype=dtype)
+        k = torch.randn(6, H, d_h, device=device, dtype=dtype)
+        v = torch.randn(6, H, d_h, device=device, dtype=dtype)
+        cu_seqlens = torch.tensor([0, 3, 6], dtype=torch.int32, device=device)
+        with pytest.raises(AssertionError, match="mutually exclusive"):
+            attn.forward_with_lse(
+                q,
+                k,
+                v,
+                key_padding_mask=torch.ones(2, 3, dtype=torch.bool, device=device),
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_kv=cu_seqlens,
+                max_seqlen_q=3,
+                max_seqlen_kv=3,
+            )
+
+
+class TestAttnImplVarlenDispatch:
+    """Attention._attn_impl_varlen_kv: the reshape/dispatch glue, isolated
+    from full Attention construction (no Linear weights / QKV proj needed --
+    this method only touches q/k/v tensors and a few scalar attributes)."""
+
+    def _make_attn_stub(self, num_heads, num_kv_heads, head_dim, supports_varlen=True):
+        stub = Attention.__new__(Attention)
+        stub.local_num_attention_heads = num_heads
+        stub.local_num_key_value_heads = num_kv_heads
+        stub.head_dim = head_dim
+        stub.supports_varlen = supports_varlen
+        stub.attn = FlashAttn4Attention(
+            num_heads=num_heads, head_dim=head_dim, num_kv_heads=num_kv_heads
+        )
+        return stub
+
+    def test_raises_when_backend_lacks_support(self):
+        stub = self._make_attn_stub(4, 4, 16, supports_varlen=False)
+        q = torch.randn(2, 3, 4 * 16)
+        k = torch.randn(5, 4 * 16)
+        v = torch.randn(5, 4 * 16)
+        cu_seqlens_kv = torch.tensor([0, 2, 5], dtype=torch.int32)
+        with pytest.raises(ValueError, match="does not support varlen"):
+            stub._attn_impl_varlen_kv(q, k, v, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_kv=3)
+
+    def test_sequence_parallel_wrapped_backend_does_not_reject_varlen_cleanly(self):
+        """Pre-fix behavior of a wrapper with supports_varlen forced True."""
+
+        class _FakeSeqParallelWrapper:
+            world_size = 4
+
+            def forward(self, q, k, v, **kwargs):
+                if q.shape[2] % self.world_size != 0:
+                    raise ValueError("num_heads not divisible by world_size")
+                return q
+
+        stub = Attention.__new__(Attention)
+        stub.local_num_attention_heads = 4
+        stub.local_num_key_value_heads = 4
+        stub.head_dim = 64
+        stub.supports_varlen = True
+        stub.attn = _FakeSeqParallelWrapper()
+
+        q = torch.randn(2, 3, 4 * 64)
+        k = torch.randn(5, 4 * 64)
+        v = torch.randn(5, 4 * 64)
+        cu_seqlens_kv = torch.tensor([0, 2, 5], dtype=torch.int32)
+
+        out = stub._attn_impl_varlen_kv(q, k, v, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_kv=3)
+        assert out.shape == (2, 3, 4 * 64)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="FA4 requires CUDA")
+    @pytest.mark.skipif(not FA4_AVAILABLE, reason="FA4 kernel not available")
+    def test_uniform_q_ragged_kv_roundtrip_shape(self):
+        device = "cuda"
+        dtype = torch.bfloat16
+        H, d_h = 4, 32
+        stub = self._make_attn_stub(H, H, d_h)
+
+        B, S_q = 2, 6
+        kv_lens = [3, 9]
+        q = torch.randn(B, S_q, H * d_h, device=device, dtype=dtype)
+        k = torch.cat([torch.randn(n, H * d_h, device=device, dtype=dtype) for n in kv_lens], dim=0)
+        v = torch.cat([torch.randn(n, H * d_h, device=device, dtype=dtype) for n in kv_lens], dim=0)
+        cu_seqlens_kv = torch.nn.functional.pad(
+            torch.cumsum(torch.tensor(kv_lens, dtype=torch.int32, device=device), dim=0), (1, 0)
+        ).to(torch.int32)
+
+        out = stub._attn_impl_varlen_kv(
+            q, k, v, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_kv=max(kv_lens)
+        )
+        assert out.shape == (B, S_q, H * d_h)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="FA4 requires CUDA")
+    @pytest.mark.skipif(not FA4_AVAILABLE, reason="FA4 kernel not available")
+    def test_uniform_q_ragged_kv_matches_sdpa_reference(self):
+        device = "cuda"
+        dtype = torch.bfloat16
+        H, d_h = 4, 32
+        stub = self._make_attn_stub(H, H, d_h)
+
+        B, S_q = 2, 6
+        kv_lens = [3, 9]
+        torch.manual_seed(2)
+        q_bhsd = torch.randn(B, H, S_q, d_h, device=device, dtype=dtype)
+        k_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        v_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        ref = _sdpa_reference(q_bhsd, k_list, v_list, stub.attn.scale)  # [B, H, S_q, d_h]
+
+        q = q_bhsd.transpose(1, 2).reshape(B, S_q, H * d_h)
+        k = torch.cat(
+            [kk.transpose(0, 1).reshape(n, H * d_h) for kk, n in zip(k_list, kv_lens)], dim=0
+        )
+        v = torch.cat(
+            [vv.transpose(0, 1).reshape(n, H * d_h) for vv, n in zip(v_list, kv_lens)], dim=0
+        )
+        cu_seqlens_kv = torch.nn.functional.pad(
+            torch.cumsum(torch.tensor(kv_lens, dtype=torch.int32, device=device), dim=0), (1, 0)
+        ).to(torch.int32)
+
+        out = stub._attn_impl_varlen_kv(
+            q, k, v, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_kv=max(kv_lens)
+        )
+        out_bhsd = out.reshape(B, S_q, H, d_h).transpose(1, 2)
+        torch.testing.assert_close(out_bhsd, ref, rtol=2e-2, atol=2e-2)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unittest/_torch/visual_gen/test_wan_transformer.py b/tests/unittest/_torch/visual_gen/test_wan_transformer.py
index b3132939f3d9..88f936bec1ed 100644
--- a/tests/unittest/_torch/visual_gen/test_wan_transformer.py
+++ b/tests/unittest/_torch/visual_gen/test_wan_transformer.py
@@ -37,13 +37,25 @@
 from diffusers import WanTransformer3DModel as HFWanTransformer3DModel
 
 from tensorrt_llm._torch.modules.linear import Linear
+from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd
 from tensorrt_llm._torch.visual_gen.config import (
     DiffusionModelConfig,
     DiffusionPipelineConfig,
     VisualGenArgs,
 )
-from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanTransformer3DModel
+from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import (
+    WanBlock,
+    WanTransformer3DModel,
+)
+from tensorrt_llm._torch.visual_gen.modules.attention import Attention
 from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.visual_gen.args import AttentionConfig
+
+FA4_AVAILABLE = _flash_attn_fwd is not None
+fa4_cuda_only = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="FA4 requires CUDA"),
+    pytest.mark.skipif(not FA4_AVAILABLE, reason="FA4 kernel not available"),
+]
 
 
 @pytest.fixture(autouse=True, scope="module")
@@ -342,6 +354,88 @@ def test_allclose_to_hf(self):
         torch.testing.assert_close(trt_out, hf_out, atol=0.4, rtol=0.4)
 
 
+def _make_fa4_model_config(config_dict: dict) -> DiffusionModelConfig:
+    return DiffusionModelConfig(
+        pretrained_config=SimpleNamespace(**config_dict),
+        quant_config=QuantConfig(),
+        quant_config_dict=None,
+        dynamic_weight_quant=False,
+        force_dynamic_quantization=False,
+        skip_create_weights_in_init=False,
+        attention=AttentionConfig(backend="FA4"),
+    )
+
+
+@pytest.mark.integration
+@pytest.mark.wan_t2v
+class TestWanBlockVarlenCrossAttn:
+    """Calls attn2 directly; ragged cu_seqlens output vs. a masked-padded oracle."""
+
+    pytestmark = fa4_cuda_only
+    DEVICE = "cuda"
+    DTYPE = torch.bfloat16
+
+    def test_varlen_matches_masked_padded_oracle(self):
+        torch.manual_seed(7)
+        hidden_size = 128
+        cfg = {
+            **WAN_1_3B_CONFIG,
+            "num_layers": 1,
+            "hidden_size": hidden_size,
+            "num_attention_heads": 2,
+            "attention_head_dim": 64,
+            "ffn_dim": hidden_size * 4,
+            "text_dim": 64,
+        }
+        block = (
+            WanBlock(model_config=_make_fa4_model_config(cfg), _layer_idx=0)
+            .to(self.DEVICE, dtype=self.DTYPE)
+            .eval()
+        )
+
+        B, seq_len, max_text_len = 2, 16, 16
+        text_lens = torch.tensor([5, max_text_len], dtype=torch.int32, device=self.DEVICE)
+
+        norm_x = torch.randn(B, seq_len, hidden_size, device=self.DEVICE, dtype=self.DTYPE)
+        encoder_hidden_states_text = torch.zeros(
+            B, max_text_len, hidden_size, device=self.DEVICE, dtype=self.DTYPE
+        )
+        for i, n in enumerate(text_lens.tolist()):
+            encoder_hidden_states_text[i, :n] = torch.randn(
+                n, hidden_size, device=self.DEVICE, dtype=self.DTYPE
+            )
+
+        with torch.inference_mode():
+            q, k, v = block.attn2.get_qkv(norm_x, encoder_hidden_states_text)
+            q, k = block.attn2.apply_qk_norm(q, k)
+
+            k_ragged, v_ragged, cu_seqlens_kv = Attention.pack_ragged_kv(k, v, text_lens)
+            varlen_out = block.attn2._attn_impl(
+                q,
+                k_ragged,
+                v_ragged,
+                batch_size=B,
+                seq_len=seq_len,
+                cu_seqlens_kv=cu_seqlens_kv,
+                max_seqlen_kv=int(text_lens.max().item()),
+            )
+
+            key_padding_mask = (
+                torch.arange(max_text_len, device=self.DEVICE)[None, :] < text_lens[:, None]
+            )
+            masked_padded_out = block.attn2._attn_impl(
+                q,
+                k,
+                v,
+                batch_size=B,
+                seq_len=seq_len,
+                kv_seq_len=max_text_len,
+                key_padding_mask=key_padding_mask,
+            )
+
+        torch.testing.assert_close(varlen_out, masked_padded_out, atol=2e-2, rtol=2e-2)
+
+
 # ============================================================================
 # T2V correctness test — Wan2.1-T2V-1.3B
 # ============================================================================

From 7c6b417d3bc4a61e17c1c7cbf5f87dff1b7690a8 Mon Sep 17 00:00:00 2001
From: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Date: Tue, 8 Sep 2026 13:48:08 -0700
Subject: [PATCH 2/9] trim varlen attention tests

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
---
 .../visual_gen/test_varlen_attention.py       | 56 -------------------
 1 file changed, 56 deletions(-)

diff --git a/tests/unittest/_torch/visual_gen/test_varlen_attention.py b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
index fc01b8752283..f12e191a0694 100644
--- a/tests/unittest/_torch/visual_gen/test_varlen_attention.py
+++ b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
@@ -17,7 +17,6 @@
 import pytest
 import torch
 
-from tensorrt_llm._torch.attention_backend.interface import PredefinedAttentionMask
 from tensorrt_llm._torch.visual_gen.attention_backend import (
     CuTeDSLAttention,
     TrtllmAttention,
@@ -171,19 +170,6 @@ def test_uneven_boundary_lengths(self):
         out, ref = _run_ragged_kv_vs_sdpa(attn, B, S_q, H, d_h, kv_lens, device, dtype)
         torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2)
 
-    def test_zero_length_kv_sample(self):
-        """Empty CFG branch (e.g. empty negative prompt) -- pins actual FA4 behavior."""
-        device, dtype = "cuda", torch.bfloat16
-        B, S_q, H, d_h = 2, 4, 8, 64
-        kv_lens = [0, 9]
-        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
-        torch.manual_seed(1)
-        q_bhsd = torch.randn(B, H, S_q, d_h, device=device, dtype=dtype)
-        k_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
-        v_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
-        out = _run_packed(attn, q_bhsd, k_list, v_list, S_q, H, d_h, device)
-        assert not torch.isnan(out).any()
-
     def test_all_equal_lengths(self):
         """Degenerate case: nothing to pack, every sample the same length."""
         device, dtype = "cuda", torch.bfloat16
@@ -193,22 +179,6 @@ def test_all_equal_lengths(self):
         out, ref = _run_ragged_kv_vs_sdpa(attn, B, S_q, H, d_h, kv_lens, device, dtype)
         torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2)
 
-    def test_raises_when_causal_combined_with_key_padding_mask(self):
-        device, dtype = "cuda", torch.bfloat16
-        H, d_h = 4, 32
-        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
-        q = torch.randn(2, 3, H, d_h, device=device, dtype=dtype)
-        k = torch.randn(2, 3, H, d_h, device=device, dtype=dtype)
-        v = torch.randn(2, 3, H, d_h, device=device, dtype=dtype)
-        with pytest.raises(AssertionError, match="key_padding_mask is not supported"):
-            attn.forward_with_lse(
-                q,
-                k,
-                v,
-                attention_mask=PredefinedAttentionMask.CAUSAL,
-                key_padding_mask=torch.ones(2, 3, dtype=torch.bool, device=device),
-            )
-
     def test_raises_when_key_padding_mask_combined_with_cu_seqlens(self):
         device, dtype = "cuda", torch.bfloat16
         H, d_h = 4, 32
@@ -255,32 +225,6 @@ def test_raises_when_backend_lacks_support(self):
         with pytest.raises(ValueError, match="does not support varlen"):
             stub._attn_impl_varlen_kv(q, k, v, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_kv=3)
 
-    def test_sequence_parallel_wrapped_backend_does_not_reject_varlen_cleanly(self):
-        """Pre-fix behavior of a wrapper with supports_varlen forced True."""
-
-        class _FakeSeqParallelWrapper:
-            world_size = 4
-
-            def forward(self, q, k, v, **kwargs):
-                if q.shape[2] % self.world_size != 0:
-                    raise ValueError("num_heads not divisible by world_size")
-                return q
-
-        stub = Attention.__new__(Attention)
-        stub.local_num_attention_heads = 4
-        stub.local_num_key_value_heads = 4
-        stub.head_dim = 64
-        stub.supports_varlen = True
-        stub.attn = _FakeSeqParallelWrapper()
-
-        q = torch.randn(2, 3, 4 * 64)
-        k = torch.randn(5, 4 * 64)
-        v = torch.randn(5, 4 * 64)
-        cu_seqlens_kv = torch.tensor([0, 2, 5], dtype=torch.int32)
-
-        out = stub._attn_impl_varlen_kv(q, k, v, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_kv=3)
-        assert out.shape == (2, 3, 4 * 64)
-
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="FA4 requires CUDA")
     @pytest.mark.skipif(not FA4_AVAILABLE, reason="FA4 kernel not available")
     def test_uniform_q_ragged_kv_roundtrip_shape(self):

From 2cc7e14a85a64f9171e9f57ffc0acd846e7957a0 Mon Sep 17 00:00:00 2001
From: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Date: Tue, 8 Sep 2026 20:17:56 -0700
Subject: [PATCH 3/9] avoid forced host-device sync in pack_ragged_kv, fix
 varlen cross-attn test flakiness

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
---
 tensorrt_llm/_torch/visual_gen/modules/attention.py  | 12 ++++++++----
 .../_torch/visual_gen/test_varlen_attention.py       |  4 ++--
 .../_torch/visual_gen/test_wan_transformer.py        |  5 ++++-
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py
index 1d7552339c86..6ef763b9b682 100644
--- a/tensorrt_llm/_torch/visual_gen/modules/attention.py
+++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py
@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -601,18 +601,22 @@ def _reshape_gate(gate: torch.Tensor) -> torch.Tensor:
 
     @staticmethod
     def pack_ragged_kv(
-        k: torch.Tensor, v: torch.Tensor, kv_lens: torch.Tensor
+        k: torch.Tensor, v: torch.Tensor, kv_lens: List[int]
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """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 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.tolist()):
+        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 = F.pad(torch.cumsum(kv_lens, dim=0), (1, 0)).to(torch.int32)
+        cu_seqlens_kv = F.pad(
+            torch.tensor(kv_lens, dtype=torch.int32, device=k.device).cumsum(dim=0), (1, 0)
+        ).to(torch.int32)
         return k_ragged, v_ragged, cu_seqlens_kv
 
     def _attn_impl_varlen_kv(
diff --git a/tests/unittest/_torch/visual_gen/test_varlen_attention.py b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
index f12e191a0694..c9d9ff46bd07 100644
--- a/tests/unittest/_torch/visual_gen/test_varlen_attention.py
+++ b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
@@ -7,8 +7,8 @@
      ``flash_attn.cute`` kernel, packed ragged K/V vs. a per-sample SDPA
      reference. Requires CUDA + the FA4 CuTe kernel.
   2. ``Attention._attn_impl_varlen_kv`` -- the dispatch/reshape glue in
-     ``modules/attention.py`` that builds ``cu_seqlens_q`` and reshapes Q/K/V
-     around the backend call. Requires CUDA (delegates to layer 1).
+     ``modules/attention.py`` that keeps Q padded and reshapes K/V around
+     the backend call.
 
 This capability is not wired into any model yet; these tests exercise the
 backend and dispatch layer directly.
diff --git a/tests/unittest/_torch/visual_gen/test_wan_transformer.py b/tests/unittest/_torch/visual_gen/test_wan_transformer.py
index 88f936bec1ed..60c57c9d8edc 100644
--- a/tests/unittest/_torch/visual_gen/test_wan_transformer.py
+++ b/tests/unittest/_torch/visual_gen/test_wan_transformer.py
@@ -392,6 +392,9 @@ def test_varlen_matches_masked_padded_oracle(self):
             .to(self.DEVICE, dtype=self.DTYPE)
             .eval()
         )
+        with torch.no_grad():
+            for p in block.attn2.parameters():
+                torch.nn.init.normal_(p, std=0.02)
 
         B, seq_len, max_text_len = 2, 16, 16
         text_lens = torch.tensor([5, max_text_len], dtype=torch.int32, device=self.DEVICE)
@@ -409,7 +412,7 @@ def test_varlen_matches_masked_padded_oracle(self):
             q, k, v = block.attn2.get_qkv(norm_x, encoder_hidden_states_text)
             q, k = block.attn2.apply_qk_norm(q, k)
 
-            k_ragged, v_ragged, cu_seqlens_kv = Attention.pack_ragged_kv(k, v, text_lens)
+            k_ragged, v_ragged, cu_seqlens_kv = Attention.pack_ragged_kv(k, v, text_lens.tolist())
             varlen_out = block.attn2._attn_impl(
                 q,
                 k_ragged,

From 4813317b353e59c5143a1bf8921ef2209accc462 Mon Sep 17 00:00:00 2001
From: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Date: Wed, 9 Sep 2026 09:33:53 -0700
Subject: [PATCH 4/9] padded q, varlen k/v for better perf

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
---
 .../attention_backend/flash_attn4.py          | 18 ++---
 .../_torch/visual_gen/modules/attention.py    | 14 +---
 .../visual_gen/test_varlen_attention.py       | 77 +++++++++++++++++++
 3 files changed, 89 insertions(+), 20 deletions(-)

diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py
index f9e033a2978e..02740bb6a63d 100644
--- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py
+++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py
@@ -183,12 +183,11 @@ def forward_with_lse(
         Forward pass returning both output and log-sum-exp (LSE).
 
         Returns:
-            output: [batch_size, seq_len, num_heads, head_dim], or
-                [total_q_tokens, num_heads, head_dim] in the varlen path
-                (cu_seqlens_kv set).
-            lse:    [batch_size, num_heads, seq_len] in the padded path, or
-                [num_heads, total_q_tokens] in the varlen path. Callers doing
-                Attention2D/Ring LSE-based combination need to handle both shapes.
+            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)
 
@@ -197,9 +196,10 @@ def forward_with_lse(
                 "cu_seqlens_kv (ragged varlen) and key_padding_mask (padded+mask) "
                 "are mutually exclusive attention modes"
             )
-            assert (
-                cu_seqlens_q is not None and max_seqlen_q is not None and max_seqlen_kv is not None
-            ), "cu_seqlens_kv requires cu_seqlens_q, max_seqlen_q, and max_seqlen_kv"
+            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,
diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py
index 6ef763b9b682..ee091ccbe7e2 100644
--- a/tensorrt_llm/_torch/visual_gen/modules/attention.py
+++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py
@@ -626,10 +626,8 @@ def _attn_impl_varlen_kv(
         v: torch.Tensor,
         **kwargs,
     ) -> torch.Tensor:
-        """Ragged K/V cross-attention: Q is uniform-length [B, S, H*D], K/V arrive
-        pre-packed as [total_kv_tokens, H_kv*D]. The caller builds
-        ``cu_seqlens_kv``/``max_seqlen_kv`` from the real per-sample lengths.
-        """
+        """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 "
@@ -641,12 +639,8 @@ def _attn_impl_varlen_kv(
         max_seqlen_kv = kwargs.pop("max_seqlen_kv")
 
         batch_size, seq_len_q = q.shape[0], q.shape[1]
-        total_q = batch_size * seq_len_q
-        cu_seqlens_q = torch.arange(
-            0, total_q + seq_len_q, seq_len_q, dtype=torch.int32, device=q.device
-        )
 
-        q = q.reshape(total_q, self.local_num_attention_heads, self.head_dim)
+        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)
 
@@ -654,9 +648,7 @@ def _attn_impl_varlen_kv(
             {
                 "batch_size": batch_size,
                 "seq_len": seq_len_q,
-                "cu_seqlens_q": cu_seqlens_q,
                 "cu_seqlens_kv": cu_seqlens_kv,
-                "max_seqlen_q": seq_len_q,
                 "max_seqlen_kv": max_seqlen_kv,
             }
         )
diff --git a/tests/unittest/_torch/visual_gen/test_varlen_attention.py b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
index c9d9ff46bd07..a3c7c8453839 100644
--- a/tests/unittest/_torch/visual_gen/test_varlen_attention.py
+++ b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
@@ -200,6 +200,83 @@ def test_raises_when_key_padding_mask_combined_with_cu_seqlens(self):
             )
 
 
+class TestFA4PaddedQRaggedK:
+    """FA4 kernel combination cu_seqlens_k set + cu_seqlens_q=None (Q stays
+    padded, only K/V ragged) -- the path Attention._attn_impl_varlen_kv uses.
+    Calls FlashAttn4Attention._fwd directly."""
+
+    pytestmark = fa4_cuda_only
+
+    def test_padded_q_ragged_k_matches_sdpa_reference(self):
+        device, dtype = "cuda", torch.bfloat16
+        B, S_q, H, d_h = 2, 4, 8, 64
+        kv_lens = [3, 9]
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+
+        torch.manual_seed(4)
+        q_bhsd = torch.randn(B, H, S_q, d_h, device=device, dtype=dtype)
+        k_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        v_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        ref = _sdpa_reference(q_bhsd, k_list, v_list, attn.scale)
+
+        q_nhd = q_bhsd.transpose(1, 2).contiguous()
+        k_ragged = torch.cat([k.transpose(0, 1) for k in k_list], dim=0)
+        v_ragged = torch.cat([v.transpose(0, 1) for v in v_list], dim=0)
+        cu_seqlens_kv = torch.nn.functional.pad(
+            torch.cumsum(torch.tensor(kv_lens, dtype=torch.int32, device=device), dim=0), (1, 0)
+        ).to(torch.int32)
+
+        out, _ = attn._fwd(
+            q_nhd,
+            k_ragged,
+            v_ragged,
+            causal=False,
+            cu_seqlens_q=None,
+            cu_seqlens_k=cu_seqlens_kv,
+            max_seqlen_q=None,
+            max_seqlen_k=max(kv_lens),
+        )
+        out_bhsd = out.transpose(1, 2)
+        torch.testing.assert_close(out_bhsd, ref, rtol=2e-2, atol=2e-2)
+
+    def test_split_consistency(self):
+        """A sub-batch run alone must match its slice of a larger batch."""
+        device, dtype = "cuda", torch.bfloat16
+        B, S_q, H, d_h = 3, 4, 8, 64
+        kv_lens = [3, 9, 5]
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+
+        torch.manual_seed(5)
+        q_bhsd = torch.randn(B, H, S_q, d_h, device=device, dtype=dtype)
+        k_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        v_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+
+        def _run(q_bhsd, k_list, v_list):
+            q_nhd = q_bhsd.transpose(1, 2).contiguous()
+            k_ragged = torch.cat([k.transpose(0, 1) for k in k_list], dim=0)
+            v_ragged = torch.cat([v.transpose(0, 1) for v in v_list], dim=0)
+            lens = [k.shape[1] for k in k_list]
+            cu_seqlens_kv = torch.nn.functional.pad(
+                torch.cumsum(torch.tensor(lens, dtype=torch.int32, device=device), dim=0), (1, 0)
+            ).to(torch.int32)
+            out, _ = attn._fwd(
+                q_nhd,
+                k_ragged,
+                v_ragged,
+                causal=False,
+                cu_seqlens_q=None,
+                cu_seqlens_k=cu_seqlens_kv,
+                max_seqlen_q=None,
+                max_seqlen_k=max(lens),
+            )
+            return out
+
+        full_out = _run(q_bhsd, k_list, v_list)
+        sub_idx = [1, 2]
+        sub_out = _run(q_bhsd[sub_idx], [k_list[i] for i in sub_idx], [v_list[i] for i in sub_idx])
+        torch.testing.assert_close(sub_out, full_out[sub_idx], rtol=2e-2, atol=2e-2)
+
+
 class TestAttnImplVarlenDispatch:
     """Attention._attn_impl_varlen_kv: the reshape/dispatch glue, isolated
     from full Attention construction (no Linear weights / QKV proj needed --

From 51201151fff4eba9eab60c0aefe30c72e46c16ca Mon Sep 17 00:00:00 2001
From: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Date: Wed, 9 Sep 2026 12:37:06 -0700
Subject: [PATCH 5/9] speed up pack_ragged_kv: host-side cumsum, cache
 cu_seqlens_kv by kv_lens

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
---
 .../_torch/visual_gen/modules/attention.py        | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py
index ee091ccbe7e2..341beb9ef029 100644
--- a/tensorrt_llm/_torch/visual_gen/modules/attention.py
+++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py
@@ -1,9 +1,9 @@
+import functools
 from enum import Enum
 from typing import List, Optional, Tuple
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from tensorrt_llm.logger import logger
 from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig
@@ -599,6 +599,15 @@ def _reshape_gate(gate: torch.Tensor) -> torch.Tensor:
         else:
             return out.flatten(2)
 
+    @staticmethod
+    @functools.lru_cache(maxsize=128)
+    def _cu_seqlens_kv_cached(kv_lens: Tuple[int, ...], device: torch.device) -> torch.Tensor:
+        """Depends only on kv_lens (constant across a sampling loop), not K/V content."""
+        cu_list = [0]
+        for n in kv_lens:
+            cu_list.append(cu_list[-1] + n)
+        return torch.tensor(cu_list, dtype=torch.int32, device=device)
+
     @staticmethod
     def pack_ragged_kv(
         k: torch.Tensor, v: torch.Tensor, kv_lens: List[int]
@@ -614,9 +623,7 @@ def pack_ragged_kv(
             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 = F.pad(
-            torch.tensor(kv_lens, dtype=torch.int32, device=k.device).cumsum(dim=0), (1, 0)
-        ).to(torch.int32)
+        cu_seqlens_kv = Attention._cu_seqlens_kv_cached(tuple(kv_lens), k.device)
         return k_ragged, v_ragged, cu_seqlens_kv
 
     def _attn_impl_varlen_kv(

From 60bb3e8bd99d72c9ce80e3fc1f1de249e97f1112 Mon Sep 17 00:00:00 2001
From: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Date: Wed, 9 Sep 2026 13:39:47 -0700
Subject: [PATCH 6/9] add production-scale and cache correctness tests for
 varlen cross-attention

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
---
 .../visual_gen/test_varlen_attention.py       | 124 ++++++++++++++++++
 1 file changed, 124 insertions(+)

diff --git a/tests/unittest/_torch/visual_gen/test_varlen_attention.py b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
index a3c7c8453839..a861970117f7 100644
--- a/tests/unittest/_torch/visual_gen/test_varlen_attention.py
+++ b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
@@ -276,6 +276,130 @@ def _run(q_bhsd, k_list, v_list):
         sub_out = _run(q_bhsd[sub_idx], [k_list[i] for i in sub_idx], [v_list[i] for i in sub_idx])
         torch.testing.assert_close(sub_out, full_out[sub_idx], rtol=2e-2, atol=2e-2)
 
+    def test_min_length_matches_sdpa_reference(self):
+        """kv_lens=1 is the real minimum: an empty negative prompt still
+        tokenizes to one EOS token, never zero (verified against the actual
+        Wan tokenizer)."""
+        device, dtype = "cuda", torch.bfloat16
+        B, S_q, H, d_h = 2, 4, 8, 64
+        kv_lens = [1, 100]
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+
+        torch.manual_seed(7)
+        q_bhsd = torch.randn(B, H, S_q, d_h, device=device, dtype=dtype)
+        k_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        v_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        ref = _sdpa_reference(q_bhsd, k_list, v_list, attn.scale)
+
+        q_nhd = q_bhsd.transpose(1, 2).contiguous()
+        k_ragged = torch.cat([k.transpose(0, 1) for k in k_list], dim=0)
+        v_ragged = torch.cat([v.transpose(0, 1) for v in v_list], dim=0)
+        cu_seqlens_kv = torch.nn.functional.pad(
+            torch.cumsum(torch.tensor(kv_lens, dtype=torch.int32, device=device), dim=0), (1, 0)
+        ).to(torch.int32)
+
+        out, _ = attn._fwd(
+            q_nhd,
+            k_ragged,
+            v_ragged,
+            causal=False,
+            cu_seqlens_q=None,
+            cu_seqlens_k=cu_seqlens_kv,
+            max_seqlen_q=None,
+            max_seqlen_k=max(kv_lens),
+        )
+        out_bhsd = out.transpose(1, 2)
+        torch.testing.assert_close(out_bhsd, ref, rtol=2e-2, atol=2e-2)
+
+    def test_production_scale_matches_sdpa_reference(self):
+        """Same combination, at Wan2.2-14B-scale shapes (720p/81-frame),
+        not just small synthetic ones."""
+        device, dtype = "cuda", torch.bfloat16
+        B, S_q, H, d_h = 2, 75600, 40, 128
+        kv_lens = [500, 5]  # worst-for-padding
+        attn = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H)
+
+        torch.manual_seed(6)
+        q_bhsd = torch.randn(B, H, S_q, d_h, device=device, dtype=dtype)
+        k_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        v_list = [torch.randn(H, n, d_h, device=device, dtype=dtype) for n in kv_lens]
+        ref = _sdpa_reference(q_bhsd, k_list, v_list, attn.scale)
+
+        q_nhd = q_bhsd.transpose(1, 2).contiguous()
+        k_padded = torch.zeros(B, 512, H, d_h, device=device, dtype=dtype)
+        v_padded = torch.zeros(B, 512, H, d_h, device=device, dtype=dtype)
+        for i, n in enumerate(kv_lens):
+            k_padded[i, :n] = k_list[i].transpose(0, 1)
+            v_padded[i, :n] = v_list[i].transpose(0, 1)
+        k_ragged, v_ragged, cu_seqlens_kv = Attention.pack_ragged_kv(k_padded, v_padded, kv_lens)
+
+        out, _ = attn._fwd(
+            q_nhd,
+            k_ragged,
+            v_ragged,
+            causal=False,
+            cu_seqlens_q=None,
+            cu_seqlens_k=cu_seqlens_kv,
+            max_seqlen_q=None,
+            max_seqlen_k=max(kv_lens),
+        )
+        out_bhsd = out.transpose(1, 2)
+        torch.testing.assert_close(out_bhsd, ref, rtol=2e-2, atol=2e-2)
+
+
+class TestPackRaggedKvCache:
+    """Attention._cu_seqlens_kv_cached correctness under repeated/interleaved
+    kv_lens. No FA4/CUDA needed -- pure tensor bookkeeping, runs on CPU."""
+
+    def setup_method(self):
+        Attention._cu_seqlens_kv_cached.cache_clear()
+
+    def _check(self, k, v, kv_lens):
+        k_ragged, v_ragged, cu = Attention.pack_ragged_kv(k, v, kv_lens)
+        expected_cu = [0]
+        for n in kv_lens:
+            expected_cu.append(expected_cu[-1] + n)
+        assert cu.tolist() == expected_cu
+        assert torch.equal(k_ragged, torch.cat([k[i, :n] for i, n in enumerate(kv_lens)], dim=0))
+        assert torch.equal(v_ragged, torch.cat([v[i, :n] for i, n in enumerate(kv_lens)], dim=0))
+
+    def test_interleaved_kv_lens_hit_and_miss_both_correct(self):
+        device, dtype = "cpu", torch.float32
+        B, S, H, d_h = 2, 96, 8, 32
+        k = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+        v = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+
+        lens_a, lens_b = [50, 14], [10, 30]
+        self._check(k, v, lens_a)  # miss
+        self._check(k, v, lens_b)  # miss
+        self._check(k, v, lens_a)  # hit
+        self._check(k, v, lens_b)  # hit
+
+        info = Attention._cu_seqlens_kv_cached.cache_info()
+        assert info.misses == 2
+        assert info.hits == 2
+
+    def test_repeated_kv_lens_content_correct_on_new_k_v(self):
+        """Same kv_lens (cache hit on cu_seqlens_kv) but different K/V content
+        each call -- k_ragged/v_ragged must reflect the new content, not a
+        stale cached value."""
+        device, dtype = "cpu", torch.float32
+        B, S, H, d_h = 2, 32, 4, 16
+        kv_lens = [7, 3]
+
+        for _ in range(3):
+            k = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+            v = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+            self._check(k, v, kv_lens)
+
+    def test_min_length_entry(self):
+        """kv_lens=1 is the real minimum (empty prompt -> one EOS token)."""
+        device, dtype = "cpu", torch.float32
+        B, S, H, d_h = 2, 96, 8, 32
+        k = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+        v = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+        self._check(k, v, [1, 60])
+
 
 class TestAttnImplVarlenDispatch:
     """Attention._attn_impl_varlen_kv: the reshape/dispatch glue, isolated

From fdaac48a44d23d8a278d7224d6ff8a7a2555742a Mon Sep 17 00:00:00 2001
From: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Date: Thu, 10 Sep 2026 14:35:19 -0700
Subject: [PATCH 7/9] drop public varlen_cfg knob, cache
 cu_seqlens_kv/max_seqlen_kv

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
---
 tensorrt_llm/_torch/visual_gen/config.py      |   6 +-
 .../_torch/visual_gen/modules/attention.py    |  48 ++++++--
 tensorrt_llm/visual_gen/args.py               |   8 --
 .../visual_gen/test_varlen_attention.py       | 111 ++++++++++++++++--
 .../_torch/visual_gen/test_wan_transformer.py |   6 +-
 5 files changed, 141 insertions(+), 38 deletions(-)

diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py
index 26ad53454c9d..7982f96fee9c 100644
--- a/tensorrt_llm/_torch/visual_gen/config.py
+++ b/tensorrt_llm/_torch/visual_gen/config.py
@@ -74,7 +74,7 @@ def discover_pipeline_components(checkpoint_path: Path) -> Dict[str, Path]:
 
 
 def create_attention_metadata_state() -> Dict[str, Any]:
-    """Create model-scoped attention metadata state for TRTLLM visual-gen backend."""
+    """Model-scoped cache dict, shared across every Attention layer of one model."""
     return {"metadata_cache": {}}
 
 
@@ -628,7 +628,9 @@ def from_pretrained(
             NVFP4LinearMethod.use_tunable_quantize = True
 
         attention_metadata_state = (
-            create_attention_metadata_state() if attention_cfg.backend == "TRTLLM" else None
+            create_attention_metadata_state()
+            if attention_cfg.backend in ("TRTLLM", "FA4")
+            else None
         )
 
         device = kwargs.pop("device", "cuda")
diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py
index 341beb9ef029..51afed3b927e 100644
--- a/tensorrt_llm/_torch/visual_gen/modules/attention.py
+++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py
@@ -1,6 +1,5 @@
-import functools
 from enum import Enum
-from typing import List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -152,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)
@@ -600,22 +600,42 @@ def _reshape_gate(gate: torch.Tensor) -> torch.Tensor:
             return out.flatten(2)
 
     @staticmethod
-    @functools.lru_cache(maxsize=128)
-    def _cu_seqlens_kv_cached(kv_lens: Tuple[int, ...], device: torch.device) -> torch.Tensor:
-        """Depends only on kv_lens (constant across a sampling loop), not K/V content."""
+    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)
-        return torch.tensor(cu_list, dtype=torch.int32, device=device)
+        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]
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        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 for _attn_impl_varlen_kv.
-        kv_lens is a plain host-side list; a CUDA tensor would force a sync
-        on .tolist().
+        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):
@@ -623,8 +643,10 @@ def pack_ragged_kv(
             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 = Attention._cu_seqlens_kv_cached(tuple(kv_lens), k.device)
-        return k_ragged, v_ragged, cu_seqlens_kv
+        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,
diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py
index f65c6a39135e..f7d483389dd0 100644
--- a/tensorrt_llm/visual_gen/args.py
+++ b/tensorrt_llm/visual_gen/args.py
@@ -119,14 +119,6 @@ class AttentionConfig(StrictBaseModel):
             "skip_softmax (TRTLLM / CUTEDSL backends) or VSA (CUTEDSL backend)."
         ),
     )
-    enable_varlen_cfg: bool = Field(
-        False,
-        status="prototype",
-        description=(
-            "Pack unequal-length CFG text cross-attention via cu_seqlens instead of "
-            "padding. Requires an FA4 backend. Not yet wired through any model."
-        ),
-    )
 
     @model_validator(mode="after")
     def _validate_quant_attention_config(self) -> "AttentionConfig":
diff --git a/tests/unittest/_torch/visual_gen/test_varlen_attention.py b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
index a861970117f7..1db5ba11c199 100644
--- a/tests/unittest/_torch/visual_gen/test_varlen_attention.py
+++ b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
@@ -26,7 +26,10 @@
     FlashAttn4Attention,
     _flash_attn_fwd,
 )
-from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig
+from tensorrt_llm._torch.visual_gen.config import (
+    DiffusionModelConfig,
+    create_attention_metadata_state,
+)
 from tensorrt_llm._torch.visual_gen.modules import attention as attention_module
 from tensorrt_llm._torch.visual_gen.modules.attention import Attention
 from tensorrt_llm.visual_gen.args import AttentionConfig
@@ -331,7 +334,7 @@ def test_production_scale_matches_sdpa_reference(self):
         for i, n in enumerate(kv_lens):
             k_padded[i, :n] = k_list[i].transpose(0, 1)
             v_padded[i, :n] = v_list[i].transpose(0, 1)
-        k_ragged, v_ragged, cu_seqlens_kv = Attention.pack_ragged_kv(k_padded, v_padded, kv_lens)
+        k_ragged, v_ragged, cu_seqlens_kv, _ = Attention.pack_ragged_kv(k_padded, v_padded, kv_lens)
 
         out, _ = attn._fwd(
             q_nhd,
@@ -348,20 +351,25 @@ def test_production_scale_matches_sdpa_reference(self):
 
 
 class TestPackRaggedKvCache:
-    """Attention._cu_seqlens_kv_cached correctness under repeated/interleaved
-    kv_lens. No FA4/CUDA needed -- pure tensor bookkeeping, runs on CPU."""
+    """Attention._cu_seqlens_kv_and_max correctness under repeated/interleaved
+    kv_lens, cached in a model-scoped metadata_state dict. No FA4/CUDA
+    needed - pure tensor bookkeeping, runs on CPU."""
 
     def setup_method(self):
-        Attention._cu_seqlens_kv_cached.cache_clear()
+        self.metadata_state = {}
 
     def _check(self, k, v, kv_lens):
-        k_ragged, v_ragged, cu = Attention.pack_ragged_kv(k, v, kv_lens)
+        k_ragged, v_ragged, cu, max_seqlen_kv = Attention.pack_ragged_kv(
+            k, v, kv_lens, metadata_state=self.metadata_state
+        )
         expected_cu = [0]
         for n in kv_lens:
             expected_cu.append(expected_cu[-1] + n)
         assert cu.tolist() == expected_cu
+        assert max_seqlen_kv == max(kv_lens)
         assert torch.equal(k_ragged, torch.cat([k[i, :n] for i, n in enumerate(kv_lens)], dim=0))
         assert torch.equal(v_ragged, torch.cat([v[i, :n] for i, n in enumerate(kv_lens)], dim=0))
+        return cu
 
     def test_interleaved_kv_lens_hit_and_miss_both_correct(self):
         device, dtype = "cpu", torch.float32
@@ -370,14 +378,14 @@ def test_interleaved_kv_lens_hit_and_miss_both_correct(self):
         v = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
 
         lens_a, lens_b = [50, 14], [10, 30]
-        self._check(k, v, lens_a)  # miss
-        self._check(k, v, lens_b)  # miss
-        self._check(k, v, lens_a)  # hit
-        self._check(k, v, lens_b)  # hit
+        cu_a_miss = self._check(k, v, lens_a)  # miss
+        cu_b_miss = self._check(k, v, lens_b)  # miss
+        cu_a_hit = self._check(k, v, lens_a)  # hit
+        cu_b_hit = self._check(k, v, lens_b)  # hit
 
-        info = Attention._cu_seqlens_kv_cached.cache_info()
-        assert info.misses == 2
-        assert info.hits == 2
+        assert cu_a_hit is cu_a_miss
+        assert cu_b_hit is cu_b_miss
+        assert len(self.metadata_state["varlen_kv_cache"]) == 2
 
     def test_repeated_kv_lens_content_correct_on_new_k_v(self):
         """Same kv_lens (cache hit on cu_seqlens_kv) but different K/V content
@@ -400,6 +408,83 @@ def test_min_length_entry(self):
         v = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
         self._check(k, v, [1, 60])
 
+    def test_no_metadata_state_still_computes_correctly(self):
+        """metadata_state is optional - omitting it (no shared model scope)
+        must still return correct, uncached results every call."""
+        device, dtype = "cpu", torch.float32
+        B, S, H, d_h = 2, 32, 4, 16
+        kv_lens = [7, 3]
+        k = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+        v = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+
+        k_ragged, v_ragged, cu, max_seqlen_kv = Attention.pack_ragged_kv(k, v, kv_lens)
+        assert cu.tolist() == [0, 7, 10]
+        assert max_seqlen_kv == 7
+
+
+class TestVarlenKvCacheSharedAcrossModel:
+    """cu_seqlens_kv/max_seqlen_kv prepared once per length layout and reused
+    across layers/denoising steps via shared attention_metadata_state.
+    One DiffusionModelConfig is built once and passed to every layer, matching
+    how WanTransformer3DModel constructs blocks."""
+
+    def _make_layers(self, config, num_layers=3):
+        return [
+            Attention(hidden_size=64, num_attention_heads=4, head_dim=16, config=config)
+            for _ in range(num_layers)
+        ]
+
+    def _make_config(self):
+        return DiffusionModelConfig(
+            attention=AttentionConfig(backend="FA4"),
+            attention_metadata_state=create_attention_metadata_state(),
+        )
+
+    def test_layers_share_one_metadata_state_instance(self):
+        config = self._make_config()
+        layers = self._make_layers(config)
+        assert all(layer._metadata_state is config.attention_metadata_state for layer in layers)
+
+    def test_cu_seqlens_prepared_once_reused_across_layers_and_steps(self):
+        config = self._make_config()
+        layers = self._make_layers(config)
+        kv_lens = [7, 3]
+        device, dtype = "cpu", torch.float32
+        B, S, H, d_h = 2, 32, 4, 16
+
+        cached_cu, cached_max = None, None
+        for _step in range(4):
+            for layer in layers:
+                k = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+                v = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+                k_ragged, v_ragged, cu, max_seqlen_kv = Attention.pack_ragged_kv(
+                    k, v, kv_lens, metadata_state=layer._metadata_state
+                )
+                if cached_cu is None:
+                    cached_cu, cached_max = cu, max_seqlen_kv
+                else:
+                    assert cu is cached_cu
+                    assert max_seqlen_kv == cached_max
+                assert torch.equal(
+                    k_ragged, torch.cat([k[i, :n] for i, n in enumerate(kv_lens)], dim=0)
+                )
+
+        assert len(config.attention_metadata_state["varlen_kv_cache"]) == 1
+
+    def test_new_length_layout_gets_its_own_entry(self):
+        config = self._make_config()
+        (layer,) = self._make_layers(config, num_layers=1)
+        device, dtype = "cpu", torch.float32
+        B, S, H, d_h = 2, 32, 4, 16
+        k = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+        v = torch.randn(B, S, H, d_h, device=device, dtype=dtype)
+
+        Attention.pack_ragged_kv(k, v, [7, 3], metadata_state=layer._metadata_state)
+        Attention.pack_ragged_kv(k, v, [10, 5], metadata_state=layer._metadata_state)
+        Attention.pack_ragged_kv(k, v, [7, 3], metadata_state=layer._metadata_state)
+
+        assert len(config.attention_metadata_state["varlen_kv_cache"]) == 2
+
 
 class TestAttnImplVarlenDispatch:
     """Attention._attn_impl_varlen_kv: the reshape/dispatch glue, isolated
diff --git a/tests/unittest/_torch/visual_gen/test_wan_transformer.py b/tests/unittest/_torch/visual_gen/test_wan_transformer.py
index 60c57c9d8edc..84219724c124 100644
--- a/tests/unittest/_torch/visual_gen/test_wan_transformer.py
+++ b/tests/unittest/_torch/visual_gen/test_wan_transformer.py
@@ -412,7 +412,9 @@ def test_varlen_matches_masked_padded_oracle(self):
             q, k, v = block.attn2.get_qkv(norm_x, encoder_hidden_states_text)
             q, k = block.attn2.apply_qk_norm(q, k)
 
-            k_ragged, v_ragged, cu_seqlens_kv = Attention.pack_ragged_kv(k, v, text_lens.tolist())
+            k_ragged, v_ragged, cu_seqlens_kv, max_seqlen_kv = Attention.pack_ragged_kv(
+                k, v, text_lens.tolist()
+            )
             varlen_out = block.attn2._attn_impl(
                 q,
                 k_ragged,
@@ -420,7 +422,7 @@ def test_varlen_matches_masked_padded_oracle(self):
                 batch_size=B,
                 seq_len=seq_len,
                 cu_seqlens_kv=cu_seqlens_kv,
-                max_seqlen_kv=int(text_lens.max().item()),
+                max_seqlen_kv=max_seqlen_kv,
             )
 
             key_padding_mask = (

From ffedc89e7efa3d3b53ecfd515f225b070d741ef6 Mon Sep 17 00:00:00 2001
From: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Date: Fri, 11 Sep 2026 09:19:06 -0700
Subject: [PATCH 8/9] add varlen attention test to CI

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
---
 tests/integration/test_lists/test-db/l0_b200.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml
index 4a217b4294c5..c1bb3a5b4793 100644
--- a/tests/integration/test_lists/test-db/l0_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_b200.yml
@@ -244,6 +244,7 @@ l0_b200:
   - unittest/_torch/visual_gen/test_attention_trtllm_sage.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

From 98f85b7eb0e460f7736519cb75123d3e580f6dbd Mon Sep 17 00:00:00 2001
From: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
Date: Fri, 11 Sep 2026 10:43:09 -0700
Subject: [PATCH 9/9] rename TestPackRaggedKvCache to
 TestPackRaggedKvMetadataCache

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>
---
 tests/unittest/_torch/visual_gen/test_varlen_attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unittest/_torch/visual_gen/test_varlen_attention.py b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
index 1db5ba11c199..eff68d8c2592 100644
--- a/tests/unittest/_torch/visual_gen/test_varlen_attention.py
+++ b/tests/unittest/_torch/visual_gen/test_varlen_attention.py
@@ -350,7 +350,7 @@ def test_production_scale_matches_sdpa_reference(self):
         torch.testing.assert_close(out_bhsd, ref, rtol=2e-2, atol=2e-2)
 
 
-class TestPackRaggedKvCache:
+class TestPackRaggedKvMetadataCache:
     """Attention._cu_seqlens_kv_and_max correctness under repeated/interleaved
     kv_lens, cached in a model-scoped metadata_state dict. No FA4/CUDA
     needed - pure tensor bookkeeping, runs on CPU."""