Admit MTP/EAGLE spec-decode steps and sliding-window layers into the Triton 3D flash-decoding path (B300, NVFP4)

[jinhuang12]

Summary

This PR makes the Triton unified-attention 3D flash-decoding (split-KV) path usable for two cases that currently fall back to the slower 2D path: (1) speculative-decode verify steps, where each forward pass carries query_len = 1 + num_speculative_tokens (e.g. 5) instead of 1, and (2) sliding-window attention layers at long context. It was developed and measured on NVIDIA B300 SXM6 (Blackwell, sm_103), TP=1, running nvidia/gemma-4-31B-it-NVFP4 with MTP speculative decoding (num_speculative_tokens=4). Base is vLLM v0.22.1 (commit 0decac0d96c42b49572498019f0a0e3600f50398).

The change is structurally guarded, not flag-gated: it reduces to byte-for-byte the original code for every non-speculative caller (query_len == 1) on global/non-sliding layers, so the default decode path on every other model and arch is unchanged. There is no new env var to set.

This branch is rebased onto current main. The optimization commit was developed and measured on v0.22.1; rebasing onto main required threading the USE_CAUSAL / USE_PER_SEQ_CAUSAL constexprs added by #45163 through the new shared window-segmentation helpers (compute_window_tile_range / compute_window_segments) so the non-causal sliding-window path and the window-relative 3D segmentation stay consistent. The numbers below were measured on v0.22.1 and have not been re-validated on main.

Optimizations

• Admit speculative-decode verify steps into 3D flash-decoding — always-on, guarded no-op; bit-exact. The launcher's 3D-admission test was hardcoded to max_seqlen_q > 1, which forced every spec-decode verify step (uniform query_len = 1 + num_spec, e.g. 5) onto the 2D path. It is replaced with max_seqlen_q > decode_query_len, where decode_query_len is derived from speculative_config (None-guarded to 1 when there is no spec). The three 3D softmax-segment buffers' first dim is resized from seq_threshold_3D to seq_threshold_3D * decode_query_len, with a fail-fast shape assert. When decode_query_len == 1 the admission test is identical to the original max_seqlen_q > 1, so non-spec callers are untouched.
• Window-relative 3D segmentation for sliding-window layers — always-on, guarded no-op; bit-exact. With a window-blind 3D segmentation, a sliding-window layer at long context with a 1024-token window keeps only ~1/16 of its parallel-softmax segments active (the rest tile masked-out history). This adds a window-relative segmentation (WINDOW_SEG_3D) so the NUM_SEGMENTS_PER_SEQ segments tile only the window's live tile range. A shared @triton.jit helper (compute_window_tile_range / compute_window_segments) is used by both the mainloop early-return and the reduce_segments mask so the two kernels cannot drift apart. Sliding-window layers are admitted to 3D only under spec-decode (decode_query_len > 1 and window_size[0] >= 0); global layers (window_size[0] < 0) and the mm_prefix path are excluded.

Fixed-batch latency

vllm bench latency — nvidia/gemma-4-31B-it-NVFP4, B300 SXM6, TP=1, NVFP4, --max-num-seqs 32, MTP spec-decode (num_speculative_tokens=4), output_len=150, CUDA graphs (FULL_AND_PIECEWISE) + torch.compile inductor, 5 iters per shape. A = base (vLLM v0.22.1, unpatched); B = this PR. Both arms run the same base commit 0decac0d9 with an identical workload; the only difference is the patched files. The grid sweeps input length × batch size in a single model load per arm. OTPS = output tokens/sec; TPOT = time per output token (both decode-path, from per-request metrics).

Input len	Batch	A E2E (s)	B E2E (s)	E2E Δ	A OTPS	B OTPS	A TPOT (ms)	B TPOT (ms)
2048	1	0.810	0.739	−8.7%	254	290	3.93	3.44
2048	8	1.642	1.443	−12.1%	1790	2051	4.47	3.90
8192	1	1.887	1.218	−35.4%	123	269	8.11	3.72
8192	8	7.353	6.763	−8.0%	406	494	19.71	16.18
27000	1	6.309	4.734	−25.0%	74.5	320.6	13.42	3.12
27000	8	37.962	35.596	−6.2%	71.1	76.9	112.60	103.99

The decode-path metrics (OTPS, TPOT) are derived from per-request timing and improve at every shape, scaling with context length — largest at input_len=27000, batch=1, where the unpatched path falls back to the 2D kernel for the entire spec-verify decode (decode time 1.999 s → 0.465 s). On end-to-end latency, only the two 27000-token shapes have a delta that clears twice the baseline's per-iteration standard deviation; the shorter-context E2E deltas sit inside that ±2σ band and are not distinguishable from run-to-run variance (at 8192/bs1, for example, the baseline per-iteration standard deviation is ~0.57 s, so the ±2σ band of ~1.14 s spans the 0.67 s mean delta). Prefill latency is identical between arms at every shape (e.g. 4.263 s vs 4.261 s at 27000/bs1), consistent with the change touching only the decode path. No shape regresses on any metric.

Correctness

GSM8K (greedy, full 1319-question split), B300, vLLM v0.22.1, same config as above:

• Base 56.41% (744/1319); this PR 55.95% (738/1319); delta −0.46pp (within the 1.0pp gate).

Both changes are structurally guarded no-ops on the default path and bit-exact on the admitted spec-decode/sliding-window paths; the GSM8K delta is consistent with FP-associativity churn from the segmented-softmax reduction order rather than an algorithmic change.

After rebasing onto main, the rebased kernel was re-checked with a differential numerics test (unified_attention vs a dense PyTorch reference) over a 74-case matrix: 2D vs 3D flash-decoding, global vs sliding-window (64/128/256/1024), pure-decode vs spec-decode-verify (query_len 3/5), causal vs per-seq-causal (the USE_PER_SEQ_CAUSAL path from #45163), MHA/GQA/MQA, head_size 128/256, and softcap on the window-relative 3D path. All 74 matched the reference (max abs error consistent with bf16 rounding). The window-relative 3D segmentation branch was confirmed to actually engage on the spec-decode + sliding-window cases (verified via the reduce_segments launch). Scope of this re-check: L40S (sm_89), bf16, kernel-level vs reference only — it is not a B300/NVFP4 run and not an end-to-end GSM8K. The B300/NVFP4 GSM8K numbers above are from the pre-rebase v0.22.1 tree and have not been re-run on main.

Changes

3 files, +351/−35 (rebased onto main), Python/Triton only — no csrc/ or C++/CUDA changes, so no recompilation is required. (The +20-line delta over the original v0.22.1 commit is the USE_CAUSAL / USE_PER_SEQ_CAUSAL threading needed to integrate with #45163.)

• Attention launcher / kernels (vllm/v1/attention/ops/triton_unified_attention.py): decode_query_len kwarg + rewritten 3D-admission gate; window-relative segmentation (WINDOW_SEG_3D) in the mainloop and reduce_segments.
• Shared Triton helpers (vllm/v1/attention/ops/triton_attention_helpers.py): compute_window_tile_range / compute_window_segments used by both kernels to keep segmentation layout bit-identical.
• Attention backend builder (vllm/v1/attention/backends/triton_attn.py): derive decode_query_len from speculative_config (defaults to 1), resize the three 3D segment buffers to seq_threshold_3D * decode_query_len, plus a buffer-shape invariant assert.

AI-assisted contribution

This change was produced by an automated process driving AMMO under human review. Each change was checked for correctness (a bit-exact comparison against the 2D reference, plus a GSM8K before/after) and benchmarked before/after for latency; all numbers above were measured on real B300 hardware under CUDA graphs and torch.compile. With the optimization reducing to a structural no-op for non-speculative, non-sliding callers, the default decode path is byte-for-byte identical to base.

[moonlghtriver5-svg]

Hi @jinhuang12 — I opened #46724, which relaxes the same max_seqlen_q > 1 gate in unified_attention. Flagging it here so we don't duplicate effort and can coordinate.

The two PRs target different, non-overlapping workloads:

• Admit MTP/EAGLE spec-decode steps and sliding-window layers into the Triton 3D flash-decoding path (B300, NVFP4) #45450 (yours) admits spec-decode verify steps (query_len = 1 + num_speculative_tokens) and sliding-window layers, with the window-relative segmentation the SWA case needs.
• [Attention] Occupancy-gated 3D segmented decode for multi-query (diffusion-LM canvas) over long KV #46724 (mine) admits the diffusion-LM canvas — a bidirectional multi-query decode where max_seqlen_q ≈ 24–256 every step — over a long KV, keyed to 2D-grid occupancy (total_num_q_blocks * num_kv_heads < CTAS_PER_SM · SM_count) plus a max_seqlen_k ≥ 8192 floor, so it fires only when the 2D grid leaves SMs idle over a long context.

We also independently converged on the same design choices: both are structurally guarded with no env flag, both reduce to byte-for-byte the original path for ordinary query_len == 1 decode, and both only enlarge the segment-softmax scratch for the multi-query case (mine allocates the larger buffers lazily, so 1-query-only models allocate nothing extra).

Because my branch only flips use_3d False→True after the existing gate, the two are additive — the overlap is essentially one line in the gate, plus sizing the segment buffers for whichever multi-query row count is larger. They can land independently; I'm also happy to rebase on top of #45450 once it merges, or to coordinate on a single combined gate that covers both 1 + num_speculative_tokens and the diffusion canvas — whatever you and the maintainers prefer.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jinhuang12.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork