[BlivionIaG]

Purpose

This is follow-up 3 from the exploration plan in #28: mixed and MTP-verify batches.

• Stacked on perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28. The base is perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28's branch, so the diff shows only this change. Merge perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28 first, then retarget this PR to rdna_extras.
• Draft: not run on gfx1030 yet. It was verified offline (see Test Result). The Test Plan is for the agent with V620 access.

Problem

RdnaAttentionImpl.forward only used the decode kernel when max_query_len <= 1. Any other batch went entirely through a prefill kernel:

• MTP verify rows (q_len = 1 + k; every step when MTP is on). Each request got one CTA per head group. That CTA's 8-64-row tile held only 1 + k valid rows, and it walked the whole context serially.
• Decode rows of mixed batches (chunked prefill under load). Same problem: one serial walk per decode row.
• Oversized prefill grids. Every prefill grid is ceil(total_tokens / BR) × heads × num_seqs, and most of those CTAs exit early. The decode rows inflate both factors.
• No FULL graphs with MTP. The builder declared UNIFORM_SINGLE_TOKEN_DECODE. With MTP, vLLM logs "... is not supported with spec-decode for attention backend ... setting cudagraph_mode=PIECEWISE", so the whole model loses FULL graphs, GDN layers included.

scenario (Qwen3.8-27B shapes, D=256)	before	after
MTP=2 verify, 8 req × 3 tok at 16k, TP=4 (H 6/1)	24 CTAs, each walking 16384 keys	2304 CTAs, ≤1024 keys each (16 splits)
7 decodes at 16k + a new 100-token prompt, TP=2 (H 12/2)	120 CTAs; each decode row walks 16384 keys	1422 CTAs, ≤1024 keys each

Change

1. Kernel: per-token causal decode. fa_rdna2_decode_paged takes an optional cu_query_lens ([num_seqs + 1] int32).
   • For a request with L query tokens and KV length S, query token i attends to S − (L − 1 − i) keys. That is the same key range the causal mask gives it in the prefill kernels, and the sliding window applies per token.
   • Each CTA finds its request with a binary search over cu_query_lens.
   • Tokens past cu[-1] (graph padding) and zero-length requests come out as zeros.
   • Without cu_query_lens every token reads its own row as before, and the output is bitwise identical to perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28.
   • Coverage: the per-head D=128/256 kernels and the opt-in GQA decode kernel. The fp8/int8 wrappers pass nullptr.
   • New schema: ..., Tensor(a!) out, Tensor? cu_query_lens=None) -> ().
2. Backend: decode-first split (rdna_attn.py).
   • Builder. __init__ calls _init_reorder_batch_threshold(1, supports_spec_as_decode=True), which gives a threshold of 1 + num_speculative_tokens with MTP and 1 otherwise. Hybrid models already reorder with the same threshold (from the GDN builder). Under the V1 runner, dense models are now reordered decode-first; before this they were not reordered.
   • build(). Calls split_decodes_and_prefills (CPU only, no sync) and stores:
     • num_decodes and num_decode_tokens;
     • decode_query_start_loc, a view that is safe to capture in a graph;
     • prefill_query_start_loc, rebased as query_start_loc[nd:] - nd_tok;
     • max_prefill_seq_len, taken from seq_lens_cpu_upper_bound, which is exact for prefill rows.
   • forward(). The single-token branch is untouched, so [ROCm][RDNA] Load and call hippihx V1 from RDNA_ATTN decode (opt-in) #27's hippihx hook merges as is.
     • Otherwise the decode kernel takes rows [:num_decode_tokens] and the prefill dispatch takes the rest. Both write into slices of the layer output.
     • The prefill dispatch moved unchanged into _forward_prefill. Its heuristics (short kernel below 4096 keys, split-K count, num_seqs <= 4) now see only the prefill rows.
3. Flags.
   • VLLM_FA_RDNA2_SPLIT_DECODE=0 restores the old routing, for A/B runs and as a kill switch.
   • VLLM_FA_RDNA2_VERIFY_FULL_GRAPH=1 reports UNIFORM_BATCH, so uniform verify batches replay as FULL graphs. It is opt-in until validated.
4. Benchmark. benchmarks/kernels/benchmark_fa_rdna2_split_decode.py times the forward pass with the split off and on, for verify and mixed batches, and prints the max output difference between the two.

Not in this PR:

• perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28's idea C: split-K for the GQA prefill kernel when a short chunk sits behind a long prefix.
• Letting one GQA-decode CTA take all 1 + k verify tokens of a request, which would read K/V once instead of 1 + k times.

Not a duplicate: the only open PR that touches FA-RDNA2 or RDNA_ATTN is #28. This PR builds on it and implements its plan item 3.

AI assistance: written with Claude Code. Per AGENTS.md, a human must review every changed line and run the GPU validation below before this leaves draft.

Test Plan

For the gfx1030 agent, after #28's build and kernel tests pass:

1. Build _rocm_C. The decode op's schema gained Tensor? cu_query_lens=None.
2. Tests.

   pytest -v tests/kernels/attention/test_fa_rdna2_writer_layout.py \
             tests/kernels/attention/test_fa_rdna2_shape_sweep.py
   

   New tests:
   • test_decode_multi_token_queries[*]: verify rows, short extends, a zero-length request and padded tokens; window 0 and 50.
   • test_forward_split_decode[*]: RdnaAttentionImpl.forward end to end on a mixed batch and a verify-only batch, split on and off, D=128 and 256.
   • test_builder_splits_decode_first: CPU-only builder logic.
3. Kernel A/B.

   python benchmarks/kernels/benchmark_fa_rdna2_split_decode.py        # H 6/1 and 12/2, bs 784, hybrid layout
   python benchmarks/kernels/benchmark_fa_rdna2_split_decode.py --heads <Hq/Hkv per rank> --block-size 1024   # Flash-Next MTP script
   python benchmarks/kernels/benchmark_fa_rdna2_split_decode.py --head-size 128 --heads 32/8 --block-size 16 --layout dense
   

   • Expect large wins for verify batches at long context and small batch.
   • max|diff| should be around 1e-3 (fp16 noise).
   • Look for a crossover (large batch, short context) where the split is slower, and report it. If one exists, tune the reorder threshold, or kv_splits for multi-token rows (fewer splits once num_decode_tokens × H_q fills the GPU).
   • If graph capture fails inside the benchmark, use --no-graph.
4. In-model correctness. Use scripts/serve_gfx1030_flashnext_mtp.sh with ATTN=fa MTP=2, and serve_gfx1030_full.sh.
   • Greedy outputs with VLLM_FA_RDNA2_SPLIT_DECODE=1 vs =0: identical, or diverging only after a near-tie.
   • MTP acceptance rate (the spec-decode metrics in the server log): split on vs off, and vs ATTN=triton.
     • This is the key number. rdna_attn.py records "online-softmax split-K drift" on MTP verify, and verify rows now use the split-K decode kernel.
   • gsm8k (tests/evals or lm_eval), split on vs off.
5. In-model performance.
   • MTP=2 TPOT at 8×16k/1k with c=8, the script's validated cell (73.8 tok/s with triton): split on vs off, and vs ATTN=triton.
   • ITL p50/p99 of decoding requests while another request prefills (16k/1k, c=8), split on vs off.
   • MTP=0, 1k/512, c=8: no regression expected. Pure decode steps are unchanged; mixed steps now split.
6. FULL-graph verify (opt-in). Repeat steps 4 and 5 with VLLM_FA_RDNA2_VERIFY_FULL_GRAPH=1.
   • The startup warning above should disappear.
   • Verify batches at the capture sizes (3/6/12/24 for MTP=2) should replay as FULL graphs.
   • Check for capture and replay faults (the script notes replay faults at intermediate sizes), text coherence, acceptance rate and TPOT.
   • If it is clean and faster, flip the default in a follow-up.
7. GQA decode (optional). VLLM_FA_RDNA2_GQA_DECODE=1 also handles cu_query_lens. If perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28's step 7 finds that kernel coherent, repeat step 4 with it.

Test Result

Offline (done):

• CPU emulation of the kernels. Same harness as perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28 (fibers with block barriers, wave32 shuffles, guard pages, NaN-poisoned buffers), with perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28 as the baseline. All 35 cases pass:
  • Single-token decode (no cu_query_lens): bitwise identical to perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28 in all 10 cases. D=128/256, G=4/6, bs 16/784, dense and hybrid layouts, windows 0/1/50/500.
  • Multi-token decode vs an fp64 reference: 11 cases, max_rel 5.7e-4. q_len 1/2/3, zero-length requests plus one padded tail token, windows 0/1/3/50, bs 13/16/784. Padded rows are exactly 0.
  • Split batches, run as in forward() (decode kernel, then each of the 4 prefill kernels on the rebased remainder): 14 cases with windows 0 and 32, max_rel 1.2e-3.
  • GQA decode kernel (VLLM_FA_RDNA2_GQA_DECODE=1; launch trace confirms it runs for D=256): same 35 cases pass.
  • Randomized fiber schedule: bit-identical to the sequential run.
• gfx1030 device build (clang 18): no warnings. VGPR/SGPR/scratch/occupancy of all 25 kernels are identical to perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28.
• Backend logic on CPU (torch 2.14 CPU, importing vLLM from this tree):
  • test_builder_splits_decode_first and test_metadata_carries_causal pass. Also passing: verify-only, pure-decode, prefill-first, non-causal, and V1-style padded requests and tokens.
  • RdnaAttentionImpl.forward with the kernel module swapped for torch reference kernels (same layouts and argument contracts), on 47 batches: mixed, verify-only, pure decode, prefill first, padded tokens, windows, D=128/256, G=2-4, thresholds 1/3/4/off. Outputs match the reference, kernel dispatch and arguments are as expected, and rows past num_actual_tokens stay untouched.
  • The benchmark script runs end to end against the same stubs.
• Lint:
  • ruff: no new findings (the file's two existing ones remain).
  • mypy via tools/pre_commit/mypy.py: no new errors (the existing two in rdna_attn.py and one in the test file remain).
• Merge check: merges cleanly into the current rdna_extras (perf(rocm): opt in to HIP MoE for concurrent MTP decode #26, [ROCm][RDNA] Load and call hippihx V1 from RDNA_ATTN decode (opt-in) #27).

On gfx1030: pending, to be filled in from the plan above.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results (offline only so far; GPU results pending)
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

🤖 Generated with Claude Code

https://claude.ai/code/session_018dwtAVjZdFLPVRucCFXBzY

---

Generated by Claude Code