[TRTLLM-13371][perf] LTX-2 FA4: enable split-KV heuristic (num_splits=0) for low-occupancy cross-attn

[luyiyun1021]

Description

The visual-gen FlashAttention-4 backend (flash_attn4.py) called _flash_attn_fwd(...) with the library default num_splits=1 — one CTA per (Q-tile × head). For LTX-2's small-Q / large-KV audio cross-attn (video→audio: Q=126 audio tokens, K=15360 video tokens, head_dim=64) that launches only ~num_heads CTAs: a single SM wave that leaves most of a B200's 148 SMs idle and is bound by each CTA's long K-reduction, so it does not scale with Ulysses.

This one-line change passes num_splits=0, enabling FA4's built-in split-KV (Flash-Decoding) heuristic: it splits the KV reduction across the idle SMs and combines the partials via LSE. The heuristic returns 1 split whenever the grid already saturates the SMs, so self-attn and other large-Q paths are a no-op (unchanged). The split-KV combine is LSE-weighted and algebraically equivalent to the unsplit softmax, so there is no precision impact — only float re-association noise.

Test Coverage

Kernel microbench (single attention, B200, µs, mean over 50 iters)

Real LTX-2 shapes, per Ulysses rank (uly N → 32/N heads/rank, video full seq 15360, audio 126, text 1024). cuDNN pad = audio padded >128 to leave the SM80 fallback (the alternative fix); FA4 ns=1 = current library default (no split, i.e. not passing num_splits); FA4 ns=0 = this PR (auto split-KV heuristic).

Ulysses=1 (32 heads/rank):

attn	q×k	cuDNN unpad	cuDNN pad>128	FA4 ns=1	FA4 ns=0
video_self	15360×15360	sm100 2780	—	2780	2821
video_text	15360×1024	sm100 229	—	231	231
audio_self	126×126	sm80 4.9	sm100 7.6	4.8	4.8
a2v (Q=vid)	15360×126	sm80 76	sm100 62	36.2	36.1
v2a (Q=aud)	126×15360	sm80 425	sm100 149	124.6	40.8
audio_text	126×1024	sm80 18	sm100 14	12.2	10.1

Ulysses=2 (16 heads/rank):

attn	q×k	cuDNN unpad	cuDNN pad>128	FA4 ns=1	FA4 ns=0
video_self	15360×15360	sm100 1409	—	1508	1463
video_text	7680×1024	sm100 123	—	128	129
audio_self	126×126	sm80 4.8	sm100 7.6	4.8	4.8
a2v (Q=vid)	7680×126	sm80 38	sm100 34	20.0	20.0
v2a (Q=aud)	126×15360	sm80 424	sm100 148	123.6	23.3
audio_text	63×1024	sm80 18	sm80 18	12.0	9.4

Ulysses=4 (8 heads/rank):

attn	q×k	cuDNN unpad	cuDNN pad>128	FA4 ns=1	FA4 ns=0
video_self	15360×15360	sm100 715	—	803	790
video_text	3840×1024	sm100 66	—	63	58
audio_self	128×128	sm80 4.8	sm100 7.6	4.8	4.8
a2v (Q=vid)	3840×128	sm80 21	sm100 21	13.0	13.1
v2a (Q=aud)	128×15360	sm80 423	sm100 148	123.7	16.5
audio_text	32×1024	sm80 18	sm80 18	12.0	9.3

num_splits=0 collapses the v2a bottleneck by 3.1× / 5.3× / 7.5× (124.6→40.8 / 123.6→23.3 / 123.7→16.5 µs at uly 1/2/4) — well below both cuDNN variants — modestly helps audio_text, and stays within noise (no-op) on every large-Q path (video_self/text, a2v, audio_self) where the grid already saturates the SMs. No regression.

End-to-end LTX-2

Single-stage, 768×1280, 121 frames, 40 steps, guidance 4.0, B200, min-of-4 timed runs. FA4 backend with num_splits=0 vs the current VANILLA (cuDNN SDPA) default backend:

GPUs	parallel	VANILLA (cuDNN)	FA4 num_splits=0	speedup
1	cfg1·uly1	31.76 s	30.85 s	−2.9%
2	cfg2·uly1	17.98 s	17.24 s	−4.1%
4	cfg2·uly2	12.49 s	12.24 s	−2.0%
8	cfg2·uly4	8.88 s	8.29 s	−6.6%

The gain is largest at 8-GPU Ulysses, where the v2a cross-attn's non-scaling hurts the cuDNN path most. The change flips a single existing kernel argument to the library's auto-heuristic (no new code path); correctness is covered by the existing visual-gen LTX-2 attention/Ulysses unit tests under tests/unittest/_torch/visual_gen/, and the split-KV combine is LSE-exact.

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.

[coderabbitai]

📝 Walkthrough

Walkthrough

FlashAttn4Attention._fwd now passes num_splits=0 explicitly to _flash_attn_fwd, accompanied by an inline comment describing the intended split-KV heuristic behavior for low-occupancy small-Q/large-KV cases. No other logic, inputs, outputs, or dtype handling were changed.

Changes

FlashAttention 4 split-KV parameter

Layer / File(s)	Summary
Pass num_splits=0 to _flash_attn_fwd tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py	Adds num_splits=0 argument and an explanatory comment to the _flash_attn_fwd call inside FlashAttn4Attention._fwd, overriding the default split-KV heuristic.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~2 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: enabling split-KV heuristic (num_splits=0) for FlashAttention-4 in LTX-2 cross-attention.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description comprehensively explains the issue, solution, test coverage, and performance results with detailed benchmarks and end-to-end validation.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Warning

Review ran into problems

🔥 Problems

Stopped waiting for pipeline failures after 30000ms. One of your pipelines takes longer than our 30000ms fetch window to run, so review may not consider pipeline-failure results for inline comments if any failures occurred after the fetch window. Increase the timeout if you want to wait longer or run a @coderabbit review after the pipeline has finished.

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[luyiyun1021]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55442 [ run ] triggered by Bot. Commit: 0df59a9 Link to invocation

[tensorrt-cicd]

PR_Github #55442 [ run ] completed with state SUCCESS. Commit: 0df59a9
/LLM/main/L0_MergeRequest_PR pipeline #44378 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[luyiyun1021]

/bot run --disable-fail-fast

[luyiyun1021]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55651 [ run ] triggered by Bot. Commit: 0df59a9 Link to invocation

[tensorrt-cicd]

PR_Github #55652 [ run ] triggered by Bot. Commit: 0df59a9 Link to invocation

[tensorrt-cicd]

PR_Github #55651 [ run ] completed with state ABORTED. Commit: 0df59a9

Link to invocation

[tensorrt-cicd]

PR_Github #55652 [ run ] completed with state FAILURE. Commit: 0df59a9
/LLM/main/L0_MergeRequest_PR pipeline #44561 completed with status: 'ABORTED'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[luyiyun1021]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55691 [ run ] triggered by Bot. Commit: 0df59a9 Link to invocation

[tensorrt-cicd]

PR_Github #55691 [ run ] completed with state SUCCESS. Commit: 0df59a9
/LLM/main/L0_MergeRequest_PR pipeline #44594 completed with status: 'SUCCESS'

CI Report

Link to invocation