[Attention] Skip stage2 for single-split Triton MLA decode

[BabyDrangoner]

Purpose

When grouped Triton MLA uses one KV split, its second-stage reduction has one input and is mathematically an identity. This PR lets stage1 write the final output and LSE directly and skips the stage2 launch.

The fast path is restricted to MLA shapes for which the normal performance heuristic selects one split. Multi-split attention, MHA/GQA, and long contexts where batch invariance alone forces one split keep the existing two-stage path. The existing scratch workspace remains reserved, so this reduces HBM traffic and one launch when eligible but does not reduce reserved memory capacity.

The final implementation reuses pointer, stride, and mode slots already present in the grouped kernel. This keeps the fallback Triton launch parameter list identical to the pre-patch implementation. For empty/padded rows, the direct path writes the attention-combine neutral values output=0 and LSE=-inf. The old identity stage happened to produce NaN/-inf for those rows; they are not consumed as ordinary outputs, and the new values are safe for DCP merging.

Duplicate-work check

I searched open PRs for Triton MLA, single split, direct output, and stage2 and found no implementation of this optimization. Related work such as #33529 (split-selection/query-load tuning) and #42885 (full CUDA graph capture) is complementary; neither writes single-split MLA results directly nor skips the identity stage2 reduction.

Test plan

.venv/bin/python -m pytest -q tests/kernels/attention/test_triton_decode_attention.py
.venv/bin/pre-commit run --files \
  vllm/v1/attention/ops/triton_decode_attention.py \
  vllm/v1/attention/backends/mla/triton_mla.py \
  tests/kernels/attention/test_triton_decode_attention.py
git diff --check

The tests cover BF16/FP16 output, BF16/FP8 KV cache, active-row bitwise equivalence with the two-stage oracle, empty rows, grouped/multi-split fallback, and CUDA graph replay with changed sequence lengths. The behavior test was also run before implementation and failed because stage2 was still called once.

Cross-platform validation

The RTX PRO 6000 and H20 runs are independent, peer validations using DeepSeek-V2-Lite-Chat and TRITON_MLA. They use the same 15-shape kernel set: 13 shapes for which the production heuristic selects one split (max_seq_len < 1024 on both GPUs), plus forced-single-split B32/S2048 and B64/S4096 controls.

The tested revisions must be kept explicit: the RTX PRO 6000 was rerun on the current final source, fa0bb68; the H20 run evaluated the initial direct-output revision, 39b3e04. The H20 results therefore validate the direct fast path on Hopper and identify the initial fallback-launch overhead, but they do not claim an H20 rerun of the final slot-reuse implementation.

	NVIDIA RTX PRO 6000 Blackwell Server Edition	NVIDIA H20
Architecture	sm120, 188 SMs	sm90, 78 SMs
Driver	580.159.03	535.161.08
PyTorch / Triton / CUDA	2.11.0+cu130 / 3.6.0 / 13.0	2.11.0+cu130 / 3.6.0 / 13.0
Tested direct revision	fa0bb68 (final)	39b3e04 (initial)
Fast-path baseline	true pre-patch two-stage source (6accb779)	true pre-patch kernel file matching 324f452f, ported to the installed vLLM tree
Microbenchmark repetition	3 fresh processes; 12 balanced A B C C B A blocks per process	12 balanced A B C C B A blocks

Result definitions

• Fast-path speedup = true pre-patch two-stage time / direct time; values above 1 are faster.
• Fast-path saving = true pre-patch two-stage time - direct time; positive values are faster.
• Fallback host delta = candidate two-stage host-enqueue time - true pre-patch host-enqueue time; positive values are overhead.
• Fixed-work throughput change is derived from inverse latency; it is not an online-serving throughput measurement.
• RTX PRO 6000 kernel aggregates are the median of the three per-process paired medians. H20 aggregates are from its balanced three-arm run. Absolute latency is not compared across GPUs.

For the three-arm microbenchmarks, A is true pre-patch two-stage, B is the tested revision with direct output disabled, and C is the tested revision with direct output enabled.

Correctness

Check	NVIDIA RTX PRO 6000 (fa0bb68)	NVIDIA H20 (39b3e04)
Target kernel tests	121 passed, 16 warnings	121 passed
Common 15-shape output and LSE vs true pre-patch	bitwise equal in 45/45 process-shape cases for both final arms	bitwise equal in 15/15 shapes
Natural greedy generation, 16 prompts	identical token arrays across all three arms; 512 token IDs total, 209 distinct IDs	identical token IDs and text; 471 token IDs total, 248 distinct IDs

All applicable pre-commit hooks passed on the final source, and git diff --check passed.

Kernel performance

The graph region includes the same output-zeroing work used by the production path.

13 production-eligible shapes	NVIDIA RTX PRO 6000 (fa0bb68)	NVIDIA H20 (39b3e04)
Graph speedup	1.021x-1.152x	1.020x-1.144x
Graph saving	1.667-2.365 us	1.584-3.309 us
Eager speedup	1.044x-1.651x	1.023x-1.436x
Eager saving	2.064-14.459 us	2.467-10.855 us

The two control shapes were benchmarked with one split even though the production heuristic selects multiple splits:

Forced-single-split control	NVIDIA RTX PRO 6000 (fa0bb68)	NVIDIA H20 (39b3e04)
B32/S2048, graph / eager	1.010x (+1.719 us) / 1.023x (+4.162 us)	1.005x (+1.336 us) / 1.007x (+1.841 us)
B64/S4096, graph / eager	0.998x (-0.693 us) / 1.002x (+0.711 us)	1.002x (+1.163 us) / 1.005x (+2.550 us)

For the RTX PRO 6000 B64/S4096 graph control, per-process savings ranged from -1.262 to +0.359 us, so the small median regression is not treated as a stable architecture-level effect. The control results are not uniformly beneficial, and the production guard remains conservative rather than forcing one split for long contexts.

End-to-end decode

Both workloads used DeepSeek-V2-Lite-Chat, TRITON_MLA, FULL_DECODE_ONLY, max_model_len=512, batch 32, 64 forced decode tokens, 3 warmups, and 12 timed samples per process. The RTX PRO 6000 used two fresh processes per arm in A-B-C-C-B-A order; its estimate is the median of the two process medians and the bracketed values are their range. The H20 used an independent two-stage/direct/two-stage run.

Metric	NVIDIA RTX PRO 6000 (fa0bb68)	NVIDIA H20 (39b3e04)
True pre-patch median estimate	481.402 ms [479.585, 483.220]	390.572 ms (midpoint of 389.971 and 391.172 ms)
Final two-stage median estimate	481.587 ms [480.752, 482.423]	not separately measured in E2E
Direct median estimate	476.560 ms [475.119, 478.001]	385.849 ms
Direct saving vs true pre-patch	4.842 ms	4.723 ms
Latency change vs true pre-patch	-1.006%	-1.209%
Fixed-work inverse-latency change	+1.016%	+1.224%
Pre-patch process/arm drift	+0.758%	+0.308%
Stage1 / stage2 launches	1,701 / 1,701 pre-patch and final two-stage; 1,701 / 0 direct	1,701 / 1,701 pre-patch; 1,701 / 0 direct

On the RTX PRO 6000, both mirrored halves independently favored direct output; direct versus the final-source two-stage arm was also -1.044%. Profiling measured target-kernel totals of 18.318 ms for true pre-patch, 18.406 ms for final two-stage, and 14.765 ms for direct output. The H20 direct samples were completely separated from its two baseline sample ranges. The forced-token arrays only establish identical fixed work because token ID 198 was the sole allowed token; natural-token correctness is reported separately above. Cross-GPU absolute latency differences are not used to infer the size or cause of the speedup.

Fallback launch audit

The initial revision added five direct-only runtime values to every stage1 launch. The H20 experiment found host-enqueue overhead on write_direct=False; the same issue was reproduced independently on the RTX PRO 6000. The final revision reuses existing grouped-kernel slots and restores the pre-patch runtime/constexpr signature.

Values below are fallback host deltas as defined above, reported for B32/S128, B32/S512, and B32/S2048 respectively:

Two-stage implementation	NVIDIA RTX PRO 6000	NVIDIA H20
Initial 39b3e04	+2.681 / +2.726 / +2.821 us	+2.342 / +2.246 / +1.848 us
Final fa0bb68	+0.084 / +0.208 / +0.187 us	not rerun; no final-H20 claim

On the RTX PRO 6000, final-minus-pre-patch CUDA-event deltas for the GPU-bound S512/S2048 cases were +0.013/-0.010 us. On H20, the initial revision's two-stage graph-region deltas across the production shapes were -0.048 to +0.016 us. These measurements support a host argument-packing cause for the initial overhead, rather than a fallback-kernel GPU regression. On the PRO 6000, the final source reduces the central host deltas from about 2.7-2.8 us to at most 0.21 us.

AI assistance

OpenAI Codex assisted with code analysis, implementation, validation and benchmark scripting, and drafting this PR description.

---

• The purpose of the PR is described above.
• The test plan includes the commands and workloads used.
• Test, model-output, before/after performance, and fallback-regression results are included.
• Duplicate-work and AI-assistance disclosures are included.
• Documentation changes are not required because this does not change a user-facing API or supported-model surface.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.