[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. Flattening Q bought nothing and cost real perf. 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.

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
Performance and memory vs. padded path	Measured: real op-level win, 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)

skew	padded (ms)	seqused_k (ms)	cu_seqlens (ms)	seqused_k speedup	cu_seqlens speedup
typical	1.624	0.607	0.666	2.68x	2.44x
moderate	1.632	0.724	0.781	2.25x	2.09x
worst-for-padding	1.633	0.966	1.031	1.69x	1.58x

cu_seqlens still trails seqused_k slightly: it pays for a real memcpy to pack ragged K/V, which seqused_k avoids. Both now stay on FA4's persistent-kernel path since Q is never flattened, which is what closed most of the gap between them.

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	12081.22	12020.42	12020.90	1.01x	1.01x
moderate	12085.41	12041.10	12044.16	1.00x	1.00x
worst-for-padding	12084.43	12052.15	12058.13	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. The near-1.00x timing matches expectations too: cross-attention is ~5.6% of a layer (self-attention ~76%), and neither mechanism speeds up the surrounding projections. Amdahl's law caps the best-case e2e win around 2-3%, below this sample size's noise floor.

Bottom line: meets criteria 1-3. Keeping Q padded makes cu_seqlens close to seqused_k, but the latter is still faster. Neither moves the needle e2e: 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_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]

I dont think we should add enable_varlen_cfg for the public configuration. Each model’s attention module should select the appropriate padded or packed path based on its attention semantics and backend capabilities.

[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.