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

[o-stoner]

@coderabbitai summary

Description

Adds cu_seqlens_kv based variable-length cross-attention to the FA4 VisualGen backend, so two CFG branches (conditional and unconditional) with unequal text lengths pack K/V into one flat, unpadded batch. Q stays padded (cu_seqlens_q=None): it's always uniform-length across the CFG batch, so there's nothing to pack, and FA4 only uses its persistent-kernel scheduling when cu_seqlens_q is unset. The capability lives in the shared Attention/backend layer (tensorrt_llm/_torch/visual_gen/modules/attention.py, attention_backend/flash_attn4.py). It is not wired into Wan's model or pipeline: Wan's checkpoint was trained with unmasked padded cross-attention, so enabling packing there would silently change generation behavior. A downstream model owner can adopt this into their own model's attention calls.

cu_seqlens_kv/max_seqlen_kv depend only on the sequence-length layout (constant across every layer and denoising step of one request), so they're prepared once and cached in the shared attention_metadata_state dict, the same pattern the TRTLLM backend already uses for its prepared-metadata cache, rather than rebuilt on every attention call.

This adds support for packed K/V, but with no demonstrated speed advantage over seqused_k for the current padded-input workload. seqused_k (padded K/V plus internal masking) already supports padded and masked attention on the FA4 backend today: LTX2's shipped configs (examples/visual_gen/configs/ltx2-*.yaml) already run on backend: FA4 and exercise this path for audio padding. This PR adds the packed (cu_seqlens) variant as an additional option, not the only path to masking. Its value over seqused_k is narrower: K/V memory savings, since seqused_k still allocates the full padded-size buffer while cu_seqlens allocates only the packed size.

Acceptance criteria

Criterion	Status
Unequal-length CFG branches packed without padding	Done: Attention.pack_ragged_kv + _attn_impl_varlen_kv
Outputs match padded reference within tolerance	Done: SDPA reference, a real Attention/WanBlock instance, and the padded-Q/ragged-K FA4 kernel path directly (TestFA4PaddedQRaggedK)
Explicit backend capability checks and fallback	Done: supports_varlen() per backend, raises on unsupported backend
cu_seqlens_kv/max_seqlen_kv prepared once per length layout, reused across layers/denoising steps	Done: cached in the shared attention_metadata_state dict (same pattern as the TRTLLM backend's prepared-metadata cache), not a process-global cache
Performance and memory vs. padded path	Measured: no demonstrated kernel-level advantage over seqused_k once metadata prep is excluded from both paths; real K/V memory savings; no measurable e2e effect. See below

Microbenchmark, cross-attention op only

B200, bf16, num_heads=40, head_dim=128, img_seq_len=75600 (720p/81-frame), max_sequence_length=512, num_cfg_pairs=1, padded baseline also on FA4.

For kernel perf, we exclude each arm's own prepare_metadata stage: not just pack_ragged_kv for cu_seqlens, but also forward_with_lse's internal key_padding_mask.sum(dim=1) to seqused_k conversion, which pays its own (smaller, but real) per-call cost. Both arms are reported two ways: _total (prep timed inside the loop, what a real caller pays today) and _kernel (prep precomputed once outside the loop, isolating the FA4 kernel call itself, the apples-to-apples comparison).

skew	padded (ms)	seqused_k_total (ms)	seqused_k_kernel (ms)	cu_seqlens_total (ms)	cu_seqlens_kernel (ms)
typical	1.649	0.610	0.587	0.632	0.586
moderate	1.651	0.731	0.713	0.751	0.718
worst-for-padding	1.653	0.983	0.961	1.010	0.965

Once each arm's own prepare-metadata stage is excluded, seqused_k_kernel and cu_seqlens_kernel are statistically tied (within ~1%, run-to-run noise): the FA4 kernel itself shows no measurable advantage either way, confirming the framing above. The _total numbers show cu_seqlens trailing seqused_k by a small, consistent margin, attributable to pack_ragged_kv's K/V gather/cat, which seqused_k never pays (its own prep, a mask-sum over O(B*S_kv) elements, is cheaper than pack_ragged_kv's O(total_kv_tokens*H*D) memcpy).

K/V size (analytic, not measured, since Q dominates empirical peak memory at these shapes): padded K/V is a fixed 20.0 MB; packed K/V saves 50.7-90.0% depending on skew, though at this absolute scale (tens of MB) it's negligible next to the model's own footprint, see e2e memory below.

E2E, full 40-layer Wan2.2-14B-scale transformer forward

B200, 720p/81-frame, single-span timing around the whole forward.

skew	padded (ms)	seqused_k (ms)	cu_seqlens (ms)	seqused_k speedup	cu_seqlens speedup
typical	12245.86	12170.97	12175.56	1.01x	1.01x
moderate	12254.34	12220.89	12212.50	1.00x	1.00x
worst-for-padding	12262.13	12238.19	12225.76	1.00x	1.00x

Peak memory is identical across all three arms (51090.6 MB): the 14B-param model's weights and self-attention activations dominate so completely that cross-attention K/V is noise.

Bottom line: meets all acceptance criteria. seqused_k already handles padded and masked attention on the FA4 backend today (LTX2's shipped configs already run it); this PR adds cu_seqlens as a packed alternative, not a new capability. At the kernel level, cu_seqlens shows no measurable speed advantage over seqused_k once both arms' metadata-prep stages are excluded. The only quantifiable win for cu_seqlens is K/V memory (50-90% depending on skew), real but small in absolute terms at production model scale. Neither mechanism moves e2e latency or memory, at least not for Wan: cross-attention isn't where Wan's time or memory goes at production scale.

Test Coverage

File	Test	Covers
tests/unittest/_torch/visual_gen/test_varlen_attention.py	TestFA4VarlenKv	FA4 kernel-level ragged K/V vs. per-sample SDPA reference, split-consistency, uneven boundary lengths
tests/unittest/_torch/visual_gen/test_varlen_attention.py	TestFA4PaddedQRaggedK	FA4 padded-Q/ragged-K combination (the shipped path) vs. SDPA reference, split-consistency
tests/unittest/_torch/visual_gen/test_varlen_attention.py	TestAttnImplVarlenDispatch, test_backend_without_varlen_support_defaults_false, test_supports_varlen_checked_post_wrap	Dispatch-level correctness and explicit backend capability checks and fallback
tests/unittest/_torch/visual_gen/test_varlen_attention.py	TestPackRaggedKvCache	cu_seqlens_kv/max_seqlen_kv cache correctness under repeated/interleaved kv_lens, with and without a shared metadata_state
tests/unittest/_torch/visual_gen/test_varlen_attention.py	TestVarlenKvCacheSharedAcrossModel	cu_seqlens_kv/max_seqlen_kv prepared once per length layout and reused across multiple layers and denoising steps via shared attention_metadata_state, matching how WanTransformer3DModel shares one config across all its blocks
tests/unittest/_torch/visual_gen/test_wan_transformer.py	TestWanBlockVarlenCrossAttn::test_varlen_matches_masked_padded_oracle	Same varlen invariant through a real Attention/WanBlock instance, real projections and QK-norm

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[chang-l]

For FA4, cu_seqlens_kv and max_seqlen_kv depend on the sequence lengths, so we could prepare them once per length layout and reuse them across layers/denoising steps
Could we follow the metadata preparation/cache pattern already used by the FlashInfer backend in #18174? Its batched prefill implementation uses shared attention_metadata_state to avoid rebuilding metadata and replanning on every compatible attention call. Or, with this attn_metadata refactor PR merged, it might be easier to update for FA4 backend.

[o-stoner]

Added cu_seqlens_kv/max_seqlen_kv caching keyed on the length layout, stored in the same shared attention_metadata_state pattern as #18174 (get-or-compute-store instead of rebuilding on every call). pack_ragged_kv now takes that state as an optional param and reuses the cached tensors on a hit. Covered by TestPackRaggedKvMetadataCache and TestVarlenKvCacheSharedAcrossModel, but not yet hooked up into any model's forward pass since we don't use variable-length cross-attention by default anywhere yet