Support speculative-decode verify on ROCm, and stop auto routing short-query prefill to AITER by demandal25 · Pull Request #362 · AMD-Ecosystem/flashinfer

demandal25 · 2026-09-12T03:33:56Z

Summary

Speculative decoding was half-present on ROCm: the acceptance sampler (chain_speculative_sampling) has been a real HIP kernel for a while, but the verify step it feeds raised NotImplementedError for q_len_per_req > 1, so chain/MTP speculative decode had no supported path. This wires it up, and fixes a routing bug found while measuring it — backend="auto" was sending every short-query paged prefill to AITER at 2.3-4.6x the cost of the in-tree kernel, with an empty backend_fallback_reason to explain it.

No new kernel and no C++ change. The ROCm decode wrapper already plans through get_batch_prefill_module("fa2", ...) when use_tensor_cores=True, and the ROCm causal mask is bottom-right aligned, which is exactly verify semantics.

What changed

Multi-token verify

flashinfer/rocm/decode.py — accept q_len_per_req > 1 in _plan_impl and run(): scale qo_indptr by the draft length, widen total_num_rows, and select MaskMode.CAUSAL at run time. Also hoists kv_lens_arr_host above the persistent-buffer writes (its own commit) so the new per-request KV check can run before any self mutation, and declares _plan_info / _batch_size / _q_len_per_req in __init__ the way the prefill wrappers already do.
tests/rocm/test_batch_decode_speculative.py — new. Numerics against a causal paged-prefill oracle, the cudagraph capture/replay contract, and every refusal.

Backend routing

flashinfer/rocm/arch_caps.py — _AITER_FLAT_GATHER_GATED_Q_LEN and its accessor, beside the soft-cap constant it mirrors.
flashinfer/rocm/prefill.py — _aiter_flat_gather_short_query predicate plus one elif in _auto_select_prefill_backend, threaded from the paged planner only. _flat_gather_short_query_reason is the single producer of the decline string, so the per-batch demotion recognises it by equality.
tests/rocm/test_batch_prefill_kernels.py — test_batch_prefill_auto_selects_aiter used qo_len=16, which now sits inside the gated region on gfx942. Its subject is the capability chain (layout, dtype, head dims) and the length was incidental, so it moves to 32 rather than having its assertion weakened; the new behaviour gets its own tests.

Architecture / design notes

Chain drafts only, and the failure is silent. The verify mask is causal, so drafts are one linear sequence per request. That fits vanilla speculative decoding, DeepSeek MTP, n-gram/prompt-lookup, and EAGLE in chain mode. It does not fit tree drafts (EAGLE-2, Medusa, SpecInfer): a causal mask lets a draft token attend to a sibling on another branch, so a tree returns plausible but wrong numbers rather than raising. Tree verification needs a topology mask via custom_mask on the prefill wrapper, and the acceptance step is chain-only regardless — chain_speculative_sampling takes (batch_size, num_speculate_tokens), not a tree. Documented in docs/rocm/backends.md.

The tile rule inverts under graph capture, which is the deployment that matters. Eager planning sizes the query tile from the average packed length, so at GQA 32/8 a draft of 4 is free. Under enable_cuda_graph the scheduler cannot see per-request lengths and bounds them by total_num_rows - batch_size + 1 (PrefillSplitQOKVIndptr, include/flashinfer/rocm/attention/scheduler.cuh:540), so the tile is 64 whenever (batch_size * (q_len_per_req - 1) + 1) * gqa_group_size > 16 — at batch 4 / draft 4 / GQA 4 that is 13 * 4 = 52, the 64 tile where eager takes 16. Ordinary decode is unaffected (q_len_per_req=1 gives gqa_group_size, matching eager). Upstream avoids this with a uniform_q_len plan parameter; the ROCm binding has none, so the bound is documented rather than fixed here. This corrects a claim an earlier revision of this PR stated without the caveat.

Two of the obvious knobs are inert on ROCm, and that is worth knowing before someone deletes them. BatchPrefillWithKVCachePlan accepts bool causal and forwards to PrefillPlan, whose signature has no such parameter — the plan-side flag is set for upstream parity only, and MaskMode in run() is what actually makes the query causal. Likewise max_q_len is dropped by the fa2 paged_run wrapper; only the AITER wrapper reads it. So a blanket revert-and-rerun proves nothing about either; the A/Bs below target the changes that are real.

The routing defect is the gather, not short queries. At a page size AITER cannot page natively, batch_prefill_paged_aiter.cu index_selects the whole KV cache into a contiguous buffer before attending. That copy is O(kv) against an O(q·kv) attention, so the overhead decays as 1/q — short queries pay it in full, long ones amortise it away. The defect is a short query relative to KV.

Three paths are deliberately left alone. Native page sizes have no gather and beat fa2 even at one query row (0.93x on gfx942, 0.54x on gfx950), so the gate is disarmed by passing max_q_len=None, reusing the existing softcap_kv_len disarm idiom. Ragged prefill dispatches through mha_varlen_fwd on already-contiguous KV. Decode is genuinely one query row and AITER wins 7 of 8 measured rows there. All three leave max_q_len at its None default, so they are disarmed by construction rather than by an op-keyed branch.

Because a native page size disarms the gate up front, the run-time probe that demotes "native" to flat-gather has to re-check it — that is the only place that knows the gather is happening after all, and it has its own test.

No matching raise on the explicit-backend path, unlike the soft-cap gate next to it. That one guards a wrong answer; this is a routing preference, and a caller who writes backend="aiter" is entitled to get it — not least to benchmark the thing this gate is about.

The reason string names the threshold, not the observed length. _aiter_auto_warned is keyed on the reason, so a per-batch value would add a set entry and re-warn for every distinct draft length a serving loop sees. This is the first shape-varying gate in the file; the soft-cap one gets away with embedding its parameter because that is fixed per config.

Benchmark results

bf16 causal, head_dim 128, batch 16, pinned at 7e2913df, ROCm 10.0 / torch 2.12 / amd-aiter 0.1.20, --refcheck clean on every row. The branch has since been rebased onto #363, which changes where AITER variants are resolved from but ships a byte-identical .so (−0.6% median in its own A/B), so these ratios stand.

Verify cost relative to q=1, eager (fa2, page 16, GQA 32/8), i.e. what a draft length costs over a single-token step:

arch	s_kv	q=1	q=2	q=4	q=8	q=16
gfx942	8192	1.00	1.00	1.01	1.52	1.54
gfx950	8192	1.00	1.01	1.01	1.97	1.98

Verifying 4 draft tokens is within 1-6% of verifying 1. The step at 8 is a tile boundary, and it is not a q constant: cta_tile_q is 16 at or below 16 and 64 above, keyed on q_len * gqa_group_size (include/flashinfer/rocm/utils.cuh:100). At GQA 32/8 that puts the step at 8; at 64/8 it arrives at 2. The docs state the product rule rather than "4 is free", which would be wrong for head configurations nobody has benchmarked — and, per the design note above, wrong under graph capture at any draft length.

Threshold for the routing gate — median of 3 runs, page 64, worst (smallest) aiter/fa2 ratio across GQA groups {4,8} × kv_len {512, 4096, 32768}. Above 1 means AITER is slower:

arch	q=1	q=2	q=4	q=8	q=9	q=12	q=16	q=24
gfx942	1.86	1.77	1.74	1.67	—	1.37	1.35	0.89
gfx950	1.85	1.75	1.47	1.25	0.91	0.89	0.86	0.54

So gate at or below 16 on gfx942 and 8 on gfx950 — the largest query length at which AITER loses on every measured group and context. The gfx950 boundary was re-measured across the q=9..11 band specifically, because the original sweep jumped 8→12 and left the crossover unobserved: GQA-8 flips at exactly q=9 (0.91x), so 8 is the last losing length rather than a rounded-down guess.

Median of 3 rather than a single sweep, because AITER is known to vary run to run. The run-to-run data is the reason the ratio is the quantity quoted: max drift was 57%, but both backends drift together under contention — one config moved 45% while its ratio held at 1.78 vs 1.92 — so absolute times are soft on a shared node and the ratio is not. Single-run numbers near 1.0 flipped sign between runs at q>=24, which is outside the gated region.

Test plan

Review notes

Round 7 is worth calling out because it contradicts what rounds 5 and 6 suggested. Three of its four findings were introduced by the fixes for rounds 2-4, not by the feature: the seq_lens length check came from the round-5 commit and turned a capacity-padded buffer (a normal serving pattern, previously accepted) into a hard error; the kv_len >= q_len_per_req guard read the caller's seq_lens rather than the paged metadata the kernel actually derives kv_len from (BatchPrefillPagedParams::get_kv_len → paged_kv_t::get_length), so an understating override — which the wrapper documents as legal — let a request with kv_len < q_len_per_req reach the kernel and produce a silently wrong result, the bottom-right causal mask at prefill.cuh:895 being evaluated against a KV range shorter than the query; and moving the plan-state commit past the C++ plan (the round-4 fix) left the eager AITER early-return never assigning _qo_indptr_buf, which it had always assigned before.

An earlier revision of this description, and the commit message, said that third case ended in a device-side FLASHINFER_ERROR. That was wrong and is corrected: the only two kv_len < qo_len guards in the ROCm tree are in SinglePrefillWithKVCacheDispatched (prefill.cuh:1637) and pod.cuh:168, and the batch paged path this wrapper uses passes through neither. A silent wrong answer is the stronger argument for the fix, not the weaker one.

Three suggestions were declined with evidence rather than adopted. Replacing the reason-string comparison with a structured verdict was raised four times; it is unnecessary because both sides of the comparison call the same producer function, so the edit-drift it describes cannot occur, and the alternative changes an arity relied on by eight call sites including two files this PR does not otherwise touch. Caching _device_arch was declined on measurement: 7.2 µs per call, three calls per plan(), ~22 µs against a decode step in the hundreds of µs — not worth a first-caller-freezes-the-arch hazard in shared code. And re-promoting AITER on a graph-enabled wrapper that has not captured was declined for the third time: capture is not observable from plan(), and the proposed _aiter_flat_gather_idx is not None predicate is always None on the path in question, which test_short_query_demotion_does_not_re_promote_under_cudagraph pins.

Copilot

🟡 Changes recommended

Multi-token shape validation and CUDA-graph freezing have correctness gaps, while routing replans and warnings are handled inconsistently.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds ROCm speculative-decode verification and avoids inefficient AITER routing for short paged-prefill queries.

Changes:

Enables causal multi-token tensor-core decode.
Adds architecture-specific AITER routing thresholds.
Adds coverage and documentation for routing, validation, and CUDA graphs.

File summaries

File	Description
`flashinfer/rocm/decode.py`	Implements multi-token decode planning and execution.
`flashinfer/rocm/prefill.py`	Adds short-query AITER fallback logic.
`flashinfer/rocm/arch_caps.py`	Defines per-architecture query thresholds.
`tests/rocm/test_batch_decode_speculative.py`	Tests speculative decode and CUDA graphs.
`tests/rocm/test_batch_prefill_kernels.py`	Tests paged-prefill routing.
`tests/rocm/test_prefill_decode_dispatch.py`	Tests backend-selection predicates and warnings.
`tests/rocm/test_arch_caps.py`	Tests architecture thresholds.
`tests/rocm/test_api_parity_runtime.py`	Updates multi-token API expectations.
`docs/rocm/backends.md`	Documents verification and routing behavior.

Review details

Files reviewed: 9/9 changed files
Comments generated: 4
Review effort level: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

…ites Pure move, no behaviour change. `_plan_impl` documents at its top that a rejected plan has to leave the wrapper replayable: the caller may still be replaying a captured graph, so nothing may be written to a persistent buffer until every check has passed. `kv_lens_arr_host` was computed after the cudagraph branch had already copied into `_paged_kv_indptr_buf`, `_paged_kv_last_page_len_buf` and `_paged_kv_indices_buf`, so any future validation keyed on per-request KV length could not honour that contract. The multi-token verify support that follows needs exactly such a check (kv_len >= q_len_per_req per request), so move the computation up to sit with the other host-side derivations, before the first `self` mutation. All three existing readers are further down the method and unaffected.

The acceptance sampler (chain_speculative_sampling) has been a real HIP kernel here for a while, but the verify step it feeds had no supported path: plan() and run() both raised NotImplementedError for q_len_per_req > 1, so chain/MTP speculative decode could not run. No new kernel is needed. With use_tensor_cores=True the decode wrapper already plans through get_batch_prefill_module("fa2", ...), and the ROCm causal mask is bottom-right aligned (rocm/attention/prefill.cuh:894), so the last of the q_len_per_req rows attends to the whole KV -- exactly verify semantics. Multi-token support is therefore the same adapter upstream uses: scale qo_indptr by q_len_per_req, widen total_num_rows, and select MaskMode.CAUSAL at run time. Two things worth knowing for anyone reading this later: - The `causal` argument to the ROCm plan binding is inert. BatchPrefillWithKVCachePlan (csrc/rocm/batch_prefill.cu:49) accepts it and forwards to PrefillPlan, whose signature (include/flashinfer/rocm/attention/scheduler.cuh:685) has no such parameter. It is set here for upstream parity only; MaskMode in run() is what actually makes the query causal. Do not read the flag as the mechanism and do not delete it as dead. - qo_indptr is scaled out-of-place. _get_range_buf returns a view into a module-global cache (flashinfer/utils.py:270), so an in-place scale would corrupt the range buffer for every unrelated caller in the process while still passing this path's own tests. Under cudagraph capture q_len_per_req becomes part of the frozen shape. _qo_indptr_buf is baked as arange(batch+1) in __init__ and was never rewritten by plan, which is correct only at q_len_per_req=1; the cudagraph branch now copies the scaled offsets in. Its length does not depend on q_len_per_req, so the captured pointer stays valid. A replan with a different q_len_per_req raises rather than silently reinterpreting the query rows. run() resolves q_len_per_req before the AITER branch, not merely before the tensor-core one: AITER's PA v1 kernel reads q as [batch, heads, dim] and would misread a multi-row query rather than reject it. Verified on gfx942 against a causal paged-prefill reference at q_len_per_req in {1,2,4,8}: bitwise identical, which is the expected result given both route to the same fa2 module.

The oracle is the paged-prefill wrapper with causal=True over a qo_indptr of stride q_len_per_req. Both sides reach the same fa2 module, so a correct adapter reproduces it bitwise; what these catch is a wrong mask mode, a wrong qo_indptr stride or a wrong total_num_rows, not kernel arithmetic. Both were A/B'd against the implementation: mask mode forced back to NON_CAUSAL: q_len_per_req=1 max_abs_err=0.00000 OK <- correctly unaffected q_len_per_req=2 max_abs_err=0.13281 MISMATCH q_len_per_req=4 max_abs_err=0.16211 MISMATCH q_len_per_req=8 max_abs_err=0.17236 MISMATCH qo_indptr scaling removed: q_len_per_req=1 OK, then a GPU memory fault in PersistentVariableLengthMergeStatesKernel -- the indptr promises batch_size rows while q carries batch_size * q_len_per_req. Head configurations 32/8 and 64/8 straddle the cta_tile_q step at q_len * gqa_group_size > 16 (rocm/utils.cuh:100), so both tile sizes are exercised. The LSE case uses kv_len=8192 to force the split-kv path, whose merge_indptr is sized from total_num_rows -- otherwise that change would be untested for LSE. Two cases guard failures that are silent rather than loud: - cudagraph replay asserts on the captured _qo_indptr_buf contents as well as the output. A missed write there leaves stride-1 offsets in the captured buffer and replay returns plausible numbers. - a rejected plan must leave the wrapper replayable (the contract at the top of _plan_impl). Nothing tested it before; the new per-request KV check is the newest way to trip it. One arm compares against backend="aiter" instead: the fa2-vs-fa2 comparison cannot catch a fault shared by both plan paths. test_api_parity_runtime's multi-token test asserted the old NotImplementedError. It now pins the surviving refusal -- the default wrapper is use_tensor_cores=False, so it declines rather than silently serving one token per request. 31 tests, all passing on gfx942 (MI300X, ROCm 10.0, aiter 0.1.20).

backends.md said multi-token decode raises NotImplementedError; it does not any more. Replace that with what a caller has to satisfy: use_tensor_cores=True (the tensor-core path is the batch-prefill kernel, which supplies the causal mask, and AITER decode requires use_tensor_cores=False), kv_len >= q_len_per_req per request, and one wrapper per q_len_per_req under graph capture. The cost rule is the part worth writing down, because the obvious reading of the benchmark is wrong. Measured at GQA 32/8, a draft length of 4 is free and 8 costs 1.2-1.6x on gfx942 / 1.4-2.0x on gfx950 -- which invites "4 is the limit". It is not a q_len constant: the query tile is 16 at or below 16 and 64 above, keyed on q_len * gqa_group_size (include/flashinfer/rocm/utils.cuh:100). At 32/8 that puts the step at 8; at 64/8 it arrives at 2. Documenting the product keeps the guidance true for head configurations nobody has benchmarked. The run() docstring also claimed an output shape of [batch_size, num_qo_heads, head_dim], which stopped being true the moment q carried more than one row per request.

…uto" `auto` selected AITER for every short-query paged prefill and was 2.3-4.6x slower than the in-tree kernel for it, with an empty backend_fallback_reason -- a deliberate choice, not a fallback, so nothing explained it. The cause is not the attention kernel. A page size AITER cannot page natively takes the flat gather in batch_prefill_paged_aiter.cu, which index_selects the whole KV cache into a contiguous buffer before attending. That copy is O(kv) against an O(q*kv) attention, so its cost decays as 1/q: short queries pay it in full, long ones amortise it away. The defect is a short query *relative to KV*, not a short query. Two consequences shape the fix: - Native paging is never gated. It has no gather and beats fa2 even at one query row (0.93x gfx942, 0.54x gfx950 at q=1, page 1024). The gate is disarmed by passing max_q_len=None, reusing the existing softcap_kv_len disarm idiom so the two cannot disagree about the same call. - Only paged prefill threads the argument. Ragged dispatches through mha_varlen_fwd on already-contiguous KV with no gather, single prefill is unmeasured, and decode is genuinely one query row where AITER wins 7 of 8 measured rows. All three leave max_q_len at its None default, so they are disarmed by construction rather than by an op-keyed branch. Threshold, median of 3 runs (amd-aiter 0.1.20 / ROCm 10.0, bf16 causal head_dim=128, page 64, batch 16), worst aiter/fa2 ratio across GQA groups {4,8} x kv_len {512, 4096, 32768}: q=1 q=2 q=4 q=8 q=12 q=16 q=24 gfx942 1.86 1.77 1.74 1.67 1.37 1.35 0.89 gfx950 1.85 1.75 1.47 1.25 0.89 0.86 0.54 so gate at or below 16 on gfx942 and 8 on gfx950. Ratios rather than absolute times: on a shared node both backends drift together -- one sweep moved 45% on the fastest config while its ratio held at 1.78 vs 1.92 -- so the ratio is the contention-robust quantity and the single-run version of this table could not be trusted near 1.0. No matching raise on the explicit-backend path, unlike the soft-cap gate next to it. That one guards a wrong answer; this is a routing preference, and a caller who asks for backend="aiter" is entitled to get it -- not least to benchmark the thing this gate is about. The reason string names the threshold rather than the observed length: _aiter_auto_warned is keyed on the reason, so a per-batch value would add an entry and re-warn for every distinct draft length a serving loop sees. A/B of the new routing test, gate disarmed: AssertionError: auto kept AITER at qo_len=16 on a gathering page size assert 'aiter' == 'fa2' test_batch_prefill_auto_selects_aiter used qo_len=16, which now sits inside the gated region on gfx942. Its subject is the capability chain -- layout, dtype, head dims -- and the length was incidental, so it moves to 32 rather than having its assertion weakened, and the new behaviour gets its own test.

Sits beside the soft-cap note because it is the same shape of decision made for a different reason: that one avoids a wrong answer, this one avoids a slow path. Says which three routes are deliberately left alone -- native paging, ragged, decode -- since "auto declines AITER for short queries" otherwise reads as a blanket rule, and applying it to any of those three would give up a win AITER genuinely has.

The verify mask is causal, which is right for a linear draft sequence and silently wrong for a tree: a causal mask lets a draft token attend to a sibling on a different branch, so EAGLE-2 / Medusa / SpecInfer would get plausible numbers rather than an error. That is worth stating where a caller will read it, since nothing in the API rejects it. Names what does fit (vanilla speculative decoding, DeepSeek MTP, n-gram/prompt-lookup, EAGLE in chain mode) and points tree users at custom_mask on the prefill wrapper. Also notes the acceptance step is chain-only regardless: chain_speculative_sampling takes (batch_size, num_speculate_tokens).

Found by a review pass over the preceding commits. All three are silent -- wrong numbers or uninitialised memory rather than an error -- which is why each gets a regression test that was confirmed to fail without the fix. 1. The short-query routing demotion was one-way and sticky. Every other constraint in _auto_select_prefill_backend is device- or config-constant, so resolving once was safe; max_q_len is the first per-batch input, and `if self._backend == "auto"` never re-resolves afterwards. A serving loop that verified one 4-token draft and then ran a long chunked prefill on the same wrapper stayed on fa2 for good. Mirrors decode.py's _backend_capacity_demoted, which exists for exactly this shape of problem. AssertionError: the earlier short-query demotion stuck: a 256-token prefill is still on fa2 2. The cudagraph frozen-shape guard sat inside `if q_len_per_req > 1`, so replanning a captured wrapper back down to the default 1 skipped it entirely, rewrote _qo_indptr_buf to stride 1, and left replay reading four times the rows it attends to. Now checked in both directions and from the first plan. That makes it conservative in the other direction -- a graph-enabled wrapper warmed at q=1 cannot later plan q>1 -- which is deliberate: capture cannot be observed from plan(), so the alternative is trusting the caller not to have captured yet. Upstream freezes the same way. Failed: DID NOT RAISE ValueError 3. run() inferred q_len_per_req with floor division and then only compared it to the planned value. At batch_size=3 a q of 5 rows floors to 1, matches, and torch.empty_like(q) hands back two rows the kernel never writes. The explicit-argument branch already validated exactly; the inferred one does now too. Failed: DID NOT RAISE ValueError Also from the same pass, each smaller but real: - run() dereferenced the paged buffer before the "call plan() first" check, so running a fresh wrapper raised AttributeError on NoneType instead. - min() builtin over a tensor: unhelpful on an empty batch and elementwise on a large one. Uses .min().item() under a batch_size guard. - The second-chance branch called _aiter_flat_gather_short_query twice and carried a comment describing a precondition it does not have. It is deliberately not conditioned on gather_q_len being None -- that is the only re-check a wrapper already resolved to "aiter" ever gets -- so the comment now says so rather than inviting someone to "fix" the asymmetry. - Dropped an unreachable use_tensor_cores check in run(): plan() refuses q_len_per_req > 1 without it and the plan/run equality ties them together. - run()'s `q` docstring still claimed [batch_size, num_qo_heads, head_dim]. - Trimmed the 22-line measurement table above the new arch_caps constant to its conclusion. The evidence belongs in git log, not in an always-loaded module (~/.claude/rules/code-style.md). Test changes: the cross-backend AITER oracle now covers both sides of the cta_tile_q step (q*group 16 and 64) rather than only the tile-16 case, since the fa2-vs-fa2 comparison cannot see a fault shared by both plan paths. Added a characterization test -- passing before and after -- pinning that an explicit backend="aiter" survives the gate, which docs/rocm/backends.md promises and nothing enforced. Declined one finding: passing `q_len_per_req > 1` as the inert `causal` plan argument. If the ROCm binding ever starts honouring it, that is the correct value, so being honoured would fix the call rather than silently change it. 192 decode/dispatch/arch tests and 17 prefill routing tests pass on gfx942.

test_paged_prefill_auto_demotes_to_fa2 covers the demotion that happens when the AITER batch_prefill bootstrap raises. Its inputs came from a shared helper pinned at qo_len=16 on page_size=16 -- at the gfx942 threshold, on a page size that gathers -- so the new short-query gate declined AITER first and the test asserted on the wrong reason: assert 'query length <= 16 on a page size AITER cannot page natively ...' .startswith('aiter batch_prefill kernel bootstrap failed') Same shape as the qo_len=16 case already fixed in test_batch_prefill_kernels.py, and missed for the same reason: the length is incidental to what either test is about. _paged_inputs takes a qo_len now and this caller passes 32, above both architectures' thresholds. The helper's other caller is test_paged_prefill_explicit_aiter_still_raises, which asks for backend="aiter" explicitly and is therefore untouched by an auto-only gate -- it kept passing throughout, which is independent evidence that the explicit backend really is honoured. Found by running the blast radius rather than the changed files alone.

A second review pass over the fix commit. Both of these are mine, both are silent, and both were created by the repairs themselves -- which is the argument for reviewing a review-response commit rather than assuming it. 1. The short-query re-promotion had no cudagraph guard. Undoing a demotion resets _backend to the requested value, which re-resolves to aiter and swaps _cached_module and _plan_info -- state an already-captured graph is still pointing at. The demotion at the probe site has carried a `demotable` guard for exactly this reason; the reset I added did not. Reachable despite cudagraph freezing batch size and total rows: a ragged batch holds both constant while lifting max_q_len over the threshold. The test does that rather than lengthening the batch, which the graph rejects outright. AssertionError: re-promoted under capture; the captured graph still points at the fa2 module and plan_info 2. The cudagraph _qo_indptr_buf write landed before plan() could still raise, so a rejected plan left a captured graph reading a stride it was not captured with. Same contract the first fix commit cites twice and this one broke. Deferred past every raise, along with the _q_len_per_req commit, which feeds the frozen-shape check and so must not move either. Mismatched elements: 4 / 5 (80.0%) Greatest absolute difference: 12 at index (4,) Also corrected, all from the same pass: - The documented tile rule was false at head_dim >= 256, where FA2DetermineCtaTileQ returns 64 unconditionally (rocm/utils.cuh:97) rather than applying the q*group <= 16 test. Said so in both the docstring and backends.md; a reader sizing a draft length from "4 is free" was being misled for that configuration. - The site-2 demotion reason claimed native paging "was unavailable" for page sizes that were never native candidates. It now describes the gather without inventing a probe failure. - Finished a comment left as a truncated sentence when the measurement table was moved to git log. - Dropped a redundant local alias. - The cross-backend AITER test guarded only on is_aiter_supported (arch), so a gated capability row made it error instead of skip; two new prefill tests lacked _skip_if_prefill_gated for the same reason. - test_the_gate_warns_once_across_different_query_lengths passed vacuously wherever AITER was declined for an unrelated constant reason: both calls returned that reason, matched, and added one entry while testing nothing. New coverage for the window interaction: q_len_per_req > 1 with window_left, where prefill.cuh derives the window iteration from kv_len - qo_len - window_left. That is the one place the two features meet, and an error there returns plausible logits. Declined again: passing the literal instead of `q_len_per_req > 1` as the inert `causal` plan argument. If the binding ever honours it, that is the correct value, so being honoured fixes the call rather than changing it. Declined as out of scope: max_token_per_sequence leaves qo_indptr_host unbound in the paged planner (UnboundLocalError). Pre-existing and unconditional -- it crashes with or without this branch -- so it wants its own fix, not a rider here. 212 decode/dispatch/fallback tests and 18 prefill routing tests pass on gfx942.

…path Third review pass. The round-2 fix deferred the cudagraph buffer write past every raise but moved _q_len_per_req to the end without moving the *eager* qo_indptr write with it, so the two could disagree: plan(q_len_per_req=4) succeeds. plan(q_len_per_req=2, <bad dtype>) writes stride-2 offsets into _qo_indptr_buf, then raises before committing _q_len_per_req. run() with 4*batch rows infers 4, matches the stale planned 4, and the kernel walks stride-2 offsets -- wrong rows, half of `out` unwritten. Both writes now happen together, past every raise. Only the tensor-core run branch reads _qo_indptr_buf and the AITER path returns before that point, so deferring it changes nothing else. Also: - run()'s "call plan() first" guard was dead under cudagraph, where the paged buffers exist from __init__. It checks _plan_info now, which is the actual "has a plan" signal, and batch size comes from _batch_size -- the value the plan was built against -- rather than a buffer length. Both via getattr: neither attribute exists before the first plan. - _backend_short_query_demoted was derived by re-running the predicate outside the selector, so it also fired when an earlier constant constraint had declined AITER first and the gate was never consulted. It reads the selector's own reason now, via a shared marker constant. - Added the DeprecationWarning upstream emits for passing q_len_per_req to run() instead of plan(); its absence was a silent parity gap for anyone porting from CUDA. - The predicate docstring claimed natively-paged calls "must never reach here", which the post-probe re-check deliberately does. - The sticky-demotion test had no escape for an unrelated AITER decline, unlike its three siblings; the window test compiled its own module inside the test body because the fixture warmed only use_sliding_window=False; and the tile-rule citations named a line number that any edit above it invalidates -- they name FA2DetermineCtaTileQ now. Declined, with the test as the evidence: matching the probe site's `demotable` predicate here (`is_cuda_graph_enabled and _aiter_flat_gather_idx is not None`). That flag is set only when AITER *was* selected; on this path fa2 was, so it is always None and the predicate waves every re-promotion through: AssertionError: re-promoted under capture; the captured graph still points at the fa2 module and plan_info So the guard stays coarse -- a graph-enabled wrapper keeps fa2 once demoted, consistent with cudagraph freezing the rest of the shape. Declined again, unchanged: max_token_per_sequence's UnboundLocalError in the paged planner is pre-existing and unconditional, and wants its own fix. 212 decode/dispatch/fallback and 16 prefill routing tests pass on gfx942.

Fourth review pass, and the same class of defect as the previous two rounds: state written before something that can still raise. Each earlier round moved one more assignment behind the Python raises and stopped there; none of them accounted for `_cached_module.plan()`, the C++ PrefillPlan, which can also fail -- on a workspace too small for the total_num_rows this PR multiplies up, for instance. And round 3 left the eager _qo_indptr_buf write ahead of the raises while moving _q_len_per_req behind them, so a rejected re-plan left the two disagreeing and run()'s consistency check could not see it. _batch_size, _q_len_per_req and _qo_indptr_buf now commit together in the common tail, past every raise including the C++ one. The AITER branch returns before that tail, so it commits the same run()-visible state itself -- without that, run() would fail on a missing _batch_size for every AITER decode. Fixing it at the class level rather than one assignment at a time is the point: the last three rounds each patched the instance in front of them. Also from this pass: - _FLAT_GATHER_REASON was declared as the shared marker between the two decline sites and the consumer, but neither site interpolated it -- they hard-coded the same literal, so the coupling held by luck and a reword would have silently disabled the per-batch demotion. Both f-strings use the constant now. - Restored upstream's non-tensor-core message. Dropping it in round 1 as unreachable was wrong in effect: the equality check that replaced it tells the caller to re-plan with a q_len_per_req the wrapper cannot accept. - Removed both function-local `import warnings`. The new one alone would not have mattered, but a pre-existing local import later in the same function made the name function-local throughout, so deleting only mine produced `F823 referenced before assignment` -- caught by ruff, not by me. - The cudagraph re-promotion test could pass vacuously wherever AITER was declined for an unrelated reason; it skips on that now, like its siblings. - test_run_before_plan_says_to_call_plan's docstring described a guard that no longer exists and pointed a reader at the wrong invariant. - New coverage for run()'s explicit-q_len_per_req branch: the DeprecationWarning and the exact-shape rejection. Every other run() test went through the inference path, so that branch shipped untested. - Trimmed an over-long comment and docstring back under the project cap. 232 tests pass on gfx942. gfx950 ran 211 with 0 failures against the previous commit; re-run pending for this one.

…sts claim Fifth review pass. Nothing in it was a new defect of mine; the substantive items were an unmeasured constant, tests claiming more than they checked, and one pre-existing gap this PR had been describing as if it were closed. The gfx950 threshold of 8 rested on a gap: q=9..11 was never measured, so the boundary was asserted where nothing had been run. Measured now, median of 3 on the pinned tree, page 64, worst aiter/fa2 across GQA {4,8} x kv {512,4096,32768}: q=8 q=9 q=10 q=11 q=12 gfx950 1.50 0.91 0.90 0.90 0.91 The GQA-8 configurations flip at exactly q=9, so 8 stands -- confirmed rather than changed, and the band is no longer asserted on faith. Test claims corrected: - test_rejected_plan_leaves_wrapper_replayable re-planned with the same indptr/indices tensors, so the buffer writes it appeared to guard were no-ops and it passed for the wrong reason. Its docstring now scopes it to the pre-write raises and names what it does *not* cover. - test_run_before_plan_says_to_call_plan built a plain wrapper while its docstring explained the cudagraph case. It builds both now; only the cudagraph one can distinguish a guard on _plan_info from one on the paged buffers, which __init__ pre-allocates. - test_cudagraph_replay_matches_eager captured without warming up, so a first-call allocation into a module-global cache could land in the graph's private pool. It happened to be warm only because an earlier helper ran with the same head count. Also: - seq_lens is flattened and length-checked, as the prefill planner does. The new per-request kv_len guard reads it, so a [batch, 1] or wrong-length tensor would have validated rows the plan never schedules. - _qo_indptr_buf is set again for non-tensor-core wrappers; gating it on use_tensor_cores removed an attribute that external callers monkeypatching decode wrappers can read. - The cudagraph copy_ is skipped when q_len_per_req == 1, where __init__'s arange is already correct, and __init__'s "no need to update it in plan/run" note -- false since multi-token landed -- is fixed. - Dropped an assert that the ValueError at the top of run() now dominates and that would vanish under python -O anyway. Known gap, stated rather than papered over: the paged-KV buffers are still written before the C++ plan() that can raise, so a rejected re-plan with a new page table leaves them ahead of _plan_info. That predates this PR and making _plan_impl transactional is a refactor of a shared 450-line method, not a rider on this one. The replayability test no longer implies otherwise. Declined -- stale premise. The finding says neither producing site interpolates _FLAT_GATHER_REASON, which round 4 had already fixed: 458: _FLAT_GATHER_REASON = "flat gather" 548: f"page natively (its {_FLAT_GATHER_REASON} copies the whole " 2602: self._backend_short_query_demoted = _FLAT_GATHER_REASON in ( 2705: f"{_FLAT_GATHER_REASON}, which does not pay off at " 2455 tests pass on gfx942; gfx950 ran 212 with 0 failures on the previous commit and is re-running for this one.

The demotion site recognised the auto-selector's short-query decline by searching its reason for the marker "flat gather". The probe site's soft-cap-on-flat-gather reason spells the same phrase "flat-gather", and only that hyphen kept the two apart -- a different cause, which must not be re-evaluated per batch, was one character from being treated as one that must. _flat_gather_short_query_reason() is now the single producer and the demotion site compares for equality, so a false positive is not expressible. _aiter_flat_gather_short_query() collapses to Optional[int] (the threshold that gated, else None) since both call sites only ever read the threshold when gated. Also covers the probe-site re-check, which had no test: a native page size disarms the gate in the selector, so when the run-time probe proves native paging unavailable the gate has to be re-checked there or a short query silently takes the gather. Verified by A/B against the branch. Co-Authored-By: Claude <noreply@anthropic.com>

run() and the cudagraph frozen-shape check reached _plan_info and _q_len_per_req through getattr() defaults because __init__ never declared them. The prefill wrappers already declare _plan_info in __init__ (prefill.py:2182, :3451); following that removes the guards and makes "never planned" an explicit None rather than an absent attribute. _q_len_per_req defaults to None, not 1: the frozen-shape check treats None as "never planned", so a 1 default would reject a first plan() at q_len_per_req=2 under capture. Two accuracy fixes alongside. The comment at the late commit claimed a rejected plan leaves the previous values standing; only those three values, since the paged-KV buffers above are already written -- plan() is not transactional overall. And the run-time mask mode now keys on the planned length rather than the inferred one: equal by the check above, but the plan is what the kernel was built for. Co-Authored-By: Claude <noreply@anthropic.com>

The documented rule -- tile 16 at q_len_per_req * gqa_group_size <= 16, 64 above -- describes eager planning only, and capture is the primary speculative-decode deployment, so stating it unqualified was wrong where it matters most. Under enable_cuda_graph the scheduler cannot see per-request lengths and bounds them by total_num_rows - batch_size + 1 (PrefillSplitQOKVIndptr, include/flashinfer/rocm/attention/scheduler.cuh:540), so the tile is 64 whenever (batch_size * (q_len_per_req - 1) + 1) * gqa_group_size > 16. At batch 4 / draft 4 / GQA 4 that is 13 * 4 = 52 -- the 64 tile, where eager takes 16. Ordinary decode is unaffected: at q_len_per_req=1 the bound is gqa_group_size, which matches eager. Upstream avoids this with a uniform_q_len plan parameter; the ROCm binding has none, so documenting the bound is what is available here. Co-Authored-By: Claude <noreply@anthropic.com>

All four are mine, in code added or moved by this branch. seq_lens was newly rejected when its length exceeded batch_size. Callers size that buffer to a capacity and reuse it every step -- the same pattern paged_kv_indices already supports -- and the previous code accepted it. Slice to batch_size instead, and raise only when it is too short to cover the batch. The check existed because zero padding sinks the new kv_len guard; slicing solves that without the compat break, and also stops max() picking up a stale value left by a longer earlier step. That guard read the caller's seq_lens, but the kernel derives kv_len from indptr/last_page_len (BatchPrefillPagedParams::get_kv_len -> paged_kv_t::get_length). The wrapper documents an understating seq_lens as legal, so such an override let a request with kv_len < q_len_per_req through. The batch paged path has no kv_len < qo_len guard -- the only two in the ROCm tree are in SinglePrefillWithKVCacheDispatched (prefill.cuh:1637) and pod.cuh:168, neither on this path -- so the result was silently wrong rather than an abort: the bottom-right causal mask at prefill.cuh:895 is evaluated against a KV range shorter than the query. It now computes from the paged metadata. Moving the plan-state commit past the C++ plan left the eager AITER early-return never assigning _qo_indptr_buf, which it always did before; it is now assigned there and declared in __init__ so it is never absent. run() skipped every shape check on q when batch_size was 0: the modulo guard short-circuited and the equality check sat in the elif arm, so any row count reached paged_run. Zero batch now requires zero rows. Also drops the assert in the non-tensor-core branch, unreachable since run() gained the _plan_info ValueError above it. Co-Authored-By: Claude <noreply@anthropic.com>

The probe-site demotions logged directly, bypassing _aiter_auto_warned. The short-query one re-evaluates per batch by design, so a wrapper that degrades on every plan() warned on every plan() -- defeating the warn-once behaviour that naming the threshold rather than the observed length was specifically meant to preserve. _warn_auto_fallback_once() is now the single logging path for all four sites. The soft-cap probe-site warning had the same gap and is fixed with it: leaving one deduped and one not, twelve lines apart, is how the next edit picks the wrong one. Co-Authored-By: Claude <noreply@anthropic.com>

It passed the accepted plan's own indptr/indices/last_page_len to the rejecting plan, so the buffers were byte-identical whether or not the writes ran ahead of the raise -- the regression it exists to catch would not have failed it. The rejecting plan now carries a reversed page mapping. Two docstrings corrected to what their test actually establishes. The sliding-window case runs the same fa2 kernel on both arms, so it cannot catch a sign error in prefill.cuh's window derivation; it covers the adapter forwarding window_left and the scaled qo_indptr. And the short-query dispatch test now skips, like its three siblings, when AITER is declined for a constant constraint that answers earlier in the chain. aiter_flat_gather_gated_q_len records the shape envelope its numbers were measured over, since the threshold is a per-arch scalar applied to head dims and GQA groups nobody benchmarked. Co-Authored-By: Claude <noreply@anthropic.com>

Copilot

🟡 Changes recommended

The metadata-based KV-length safety check lacks a regression test covering a conflicting seq_lens override.

Get a fresh assessment by requesting another Copilot review.

Review details

Files reviewed: 10/10 changed files
Comments generated: 1
Review effort level: Balanced

The guard was changed to read indptr/last_page_len instead of a caller-supplied seq_lens, but every test for it passed seq_lens=None -- so reverting that line would have reintroduced the silent wrong-result path with the suite still green. Covers the disagreeing case directly: metadata reports 16 tokens per request, the override claims 512, q_len_per_req is 32. A/B against the reverted guard confirms the test fails without the fix. Also corrects the sibling docstring, which said the C++ dispatch aborts. It does not on this path -- see c0e69c1. Co-Authored-By: Claude <noreply@anthropic.com>

Copilot

🟡 Changes recommended

Documentation still misstates the guard failure mode and CUDA-graph routing behavior.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

docs/rocm/backends.md:441

This is not unconditional for CUDA-graph wrappers. After such a wrapper has planned AITER's flat-gather path, _aiter_flat_gather_idx is non-None, so the demotable check at flashinfer/rocm/prefill.py:2725-2731 deliberately keeps later short-query plans on AITER to avoid invalidating a captured graph. Document that exception (and recommend a separate wrapper) so users do not rely on the stated routing guarantee.

Files reviewed: 10/10 changed files
Comments generated: 2
Review effort level: Balanced

Copilot

🔵 Needs a closer look

The cross-backend test fails instead of skipping when AITER is unavailable on a supported GPU.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

tests/rocm/test_batch_decode_speculative.py:188

is_aiter_supported() only checks the GPU architecture, so on a supported ROCm GPU where the optional aiter package is absent or below the ABI floor this test proceeds to the explicit-AITER reference and fails instead of skipping. The other AITER kernel tests guard both conditions; include _aiter_ops_importable() here as well.

Files reviewed: 10/10 changed files
Comments generated: 0 new
Review effort level: Balanced

…e arch is_aiter_supported() checks HIP and the architecture only, so on a gfx942/gfx950 box without an importable aiter -- or one below the ABI floor -- this test fell through to the explicit backend="aiter" reference and failed where every other AITER test skips. The skip message already claimed to cover the package; now it does. Co-Authored-By: Claude <noreply@anthropic.com>

demandal25 · 2026-09-13T05:18:34Z

Suppressed comment from review 5189615732 (tests/rocm/test_batch_decode_speculative.py:188): accepted, fixed in 922196d — the guard now checks _aiter_ops_importable() too, matching the other AITER tests. Verified the claim first: is_aiter_supported checks torch.version.hip and the arch only (aiter_utils.py:23-31), so the skip message was already promising more than the guard delivered.

Copilot

🔵 Needs a closer look

Runtime and selector demotions use different reason strings, defeating warning deduplication.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

flashinfer/rocm/prefill.py:2765

This creates a second warning key for the same short-query demotion. A wrapper planned short → long → short first logs _flat_gather_short_query_reason(...) in the selector, then logs this differently worded reason in the runtime re-check, so the advertised (device, reason) warn-once behavior emits twice. Reuse the shared reason producer here so both paths deduplicate consistently.

Files reviewed: 10/10 changed files
Comments generated: 0 new
Review effort level: Balanced

The gate is reachable from two sites -- the auto selector when the page size is known to gather, and the run-time probe when a wrapper already resolved to aiter meets a short query. Each phrased its own reason, so the (device, reason) dedup saw two keys for one cause: a wrapper cycling short -> long -> short warns from the selector, then again from the probe site. The probe site now uses _flat_gather_short_query_reason() as well. The wording stays accurate there: by that point the page size does gather, whether it was never a native candidate or the probe just demoted it. This drops page_size from the message; backend_fallback_reason is the programmatic surface and the log is a log. Completes 267b250, which routed all four sites through one logger but left two of them producing different keys for the same condition. Co-Authored-By: Claude <noreply@anthropic.com>

demandal25 · 2026-09-13T05:45:41Z

Suppressed comment from review 5189634818 (flashinfer/rocm/prefill.py:2765): accepted, fixed in cb44a1e — the probe site now uses _flat_gather_short_query_reason() too, so short → long → short produces one warn-once key instead of two. Completes 267b250, which unified the logger but left two phrasings for the same condition.

Copilot

🔵 Needs a closer look

GPU-specific numerical, performance-routing, and CUDA Graph behavior warrants final human validation despite comprehensive tests.

Review details

Files reviewed: 10/10 changed files
Comments generated: 0 new
Review effort level: Balanced

Copilot

🔵 Needs a closer look

Core attention planning, backend routing, and CUDA Graph behavior require final GPU-backed human validation.

Review details

Files reviewed: 10/10 changed files
Comments generated: 0 new
Review effort level: Balanced

Copilot

🔵 Needs a closer look

GPU scheduling, masking, backend routing, and graph-capture behavior require final hardware-backed human validation.

Review details

Files reviewed: 10/10 changed files
Comments generated: 0 new
Review effort level: Balanced

#362 landed the short-query flat-gather gate in the same two places this branch changes, so `flashinfer/rocm/prefill.py` conflicted in four hunks. Both helper blocks and both new selector parameters (`allow_fp8`, `max_q_len`) are additive and are kept as-is. The two behavioural hunks needed a decision: **Re-resolution.** Each side un-sticks `plan()`'s cached backend, but for a different cause. #362 resets only a short-query demotion, via `_backend_short_query_demoted`; this branch re-resolves `auto` outright, because the poisoning case it found was a *custom-mask* demotion refusing a later fp8 plan, which the flag does not cover. Kept both: the flag still drives #362's probe-site demotion and the test that asserts it, and the general re-resolution runs after it. Both skip CUDA-graph capture, for the same reason. **`gather_q_len` was measuring the wrong set.** #362 disarmed the short-query gate when `page_size in _aiter_native_page_sizes()`, which was right before this branch and is not after it: #366 split capability from routing, and bf16 at page 16 is natively *capable* while still being *routed* to the flat gather. Left as merged, the gate would have disarmed on exactly the configuration whose gather cost it exists to avoid. It now tests `_aiter_paged_route_page_sizes`, the same predicate `softcap_kv_len` uses two lines above. Verified on gfx942 after the merge -- both probes unchanged: single aiter+fp8 : RuntimeError: AITER backend supports fp16/bf16 only ragged aiter+fp8 : NotImplementedError: ... ragged batch prefill paged aiter+fp8 page=128 : NotImplementedError: ... flat-gather route paged aiter+fp8 page=16 : NO RAISE plan 1 (custom mask) fa2 -> plan 2 (fp8) aiter Co-Authored-By: Claude <noreply@anthropic.com>

#365 landed while the previous merge was being verified. Only README.md conflicted; it is generated, so it was regenerated from arch_caps.py rather than hand-resolved, and `scripts/gen_arch_support_matrix.py --check` passes. aiter_loader.cc auto-merged cleanly -- #365's asm handle uses a fixed module name and never calls build_so_name, so this branch's fourth argument does not reach it. Also folded in the findings a `/code-review xhigh` raised against the *first* merge: **#362's probe-site test had been silently retired.** It picked its page size from `_aiter_native_page_sizes()`, which the previous merge stopped being the predicate that disarms the gate: native set : [1, 16, 1024] route set : [1024] test picks page_size = 1 does the selector disarm the gate at that page? False So the selector returned fa2 and the probe site never ran, while all four assertions still passed at the selector site. It now picks from the route set; A/B confirms it: disabling the probe-site re-check fails it, where before the fix that deletion was invisible. **fp8 at a non-routed page size with a short query raised the wrong message.** The short-query gate reaches fa2 first, so the caller was told fp8 "has no in-tree fa2 kernel ... use BatchPrefillWithPagedKVCacheWrapper" -- advice they were already following. The page-size refusal now runs first: NotImplementedError: fp8 prefill ... is not supported on the paged flat-gather route The remaining edits are small: the ragged fp8 guard moves above the jit-module split so a supplied module is covered too; the two now-unreachable `_reject_fp8_on_fa2` re-guards are dropped; `_aiter_paged_route_page_sizes` gains `functools.cache` (three calls per plan, each allocating a frozenset, on a per-step path); two tests stop mutating `global PAGE` when the helper already takes `page=`; and three comments that still justified a page size by the capability set now name the route set, which is how the retired test happened. Declined: making `_native_fp8_dtype` device-aware (flashinfer is one GPU per process), the partial `_backend` mutation before a guard raises, and the dead `scale_q/k/v` triple in `aiter_paged_run` -- all pre-existing and out of scope for a conflict resolution. Co-Authored-By: Claude <noreply@anthropic.com>

…red to gate it (#372) ## Summary The short-query section explains why ragged prefill is not gated with an argument — it runs on already-contiguous KV, so there is no gather to amortise — which reads as "no penalty". The measurement does not say that, and the sweep that sited the paged gate covers ragged too. The real reason is narrower and more useful: **gfx942 could be gated at ≤16 and is not, while gfx950 cannot be gated at all**, because one shape favours AITER from the shortest query measured onward. Follow-up to #366 and #362, and the only thing worth keeping from #367, now closed as superseded. ## What changed - **`docs/rocm/backends.md`** — six lines in the "deliberately not gated" paragraph: what the sweep shows per architecture, and which of the two gates is merely unclaimed rather than unsafe. ## Benchmark results Re-measured on merged main (`fa37f6408`) through `BatchPrefillWithRaggedKVCacheWrapper`, 9 query lengths × 5 shapes × 2 architectures, `--refcheck` clean, `auto` resolved to `aiter` in all 90 rows. Ratio is `aiter / fa2` median time back-to-back in one process, so **>1 means AITER is slower**. **gfx942 — loses in every shape at 16, first wins at 24** | bs / kv | q16 | q24 | q32 | q48 | q64 | q96 | q128 | q192 | q256 | | :--- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | | bs1 / kv1024 | 1.18 | 1.24 | 1.24 | 1.16 | 1.08 | 1.07 | 1.01 | 0.80 | 0.69 | | bs1 / kv8192 | 4.74 | 2.88 | 2.88 | 2.33 | 1.81 | 1.50 | 1.09 | 0.81 | 0.58 | | bs8 / kv1024 | 1.28 | 0.99 | 0.97 | 0.82 | 0.65 | 0.43 | 0.43 | 0.33 | 0.31 | | bs8 / kv8192 | 1.24 | 0.63 | 0.61 | 0.48 | 0.33 | 0.25 | 0.23 | 0.20 | 0.19 | | bs32 / kv2048 | 1.22 | 0.65 | 0.65 | 0.39 | 0.38 | 0.29 | 0.23 | 0.27 | 0.22 | **gfx950 — bs32/kv2048 favours AITER at every query length** | bs / kv | q16 | q24 | q32 | q48 | q64 | q96 | q128 | q192 | q256 | | :--- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | | bs1 / kv1024 | 1.29 | 1.47 | 1.26 | 1.14 | 1.05 | 1.02 | 0.90 | 0.60 | 0.63 | | bs1 / kv8192 | 5.29 | 4.28 | 3.30 | 2.02 | 1.59 | 1.17 | 0.91 | 0.55 | 0.45 | | bs8 / kv1024 | 1.11 | 0.81 | 0.71 | 0.43 | 0.38 | 0.28 | 0.25 | 0.29 | 0.15 | | bs8 / kv8192 | 1.35 | 0.71 | 0.54 | 0.31 | 0.26 | 0.18 | 0.12 | 0.24 | 0.12 | | bs32 / kv2048 | **0.58** | 0.34 | 0.34 | 0.37 | 0.21 | 0.22 | 0.17 | 0.22 | 0.16 | A gfx942-only gate at ≤16 would gain 1.18–4.74×. It is not claimed here — this PR documents the state; taking the win is a separate change with its own test surface. ## Test plan - [x] Both sweeps re-run on merged main rather than quoted from the earlier branch, and the CSVs kept outside any worktree - [x] `pre-commit run -a` - [x] `gh api /markdown` render check on the edited paragraph - [x] `/code-review medium` — the level `review-level.sh --explain` computes for a docs-only changelist No test changes: this documents a decision already implemented and covered by `test_batch_prefill_kernels.py`'s short-query tests. ## Note on two earlier drafts Worth recording, since the diff is small and the reasoning changed twice. The first draft quoted only the one cell where the sign flips and argued "a shared threshold would regress CDNA4" — which the per-arch table one screen above disproves. The second claimed ragged was ungated "for want of a crossover sweep", which Copilot correctly caught as contradicting this PR's own description: the sweep exists. Both are replaced by what the data actually shows. Co-authored-by: Claude <noreply@anthropic.com>

Four corrections from review, one of which was load-bearing for this PR's own conclusion. The perf gates decided before the amd-aiter importability check, and the selector returns as soon as a reason is set -- so on a box with no amd-aiter installed, a short query was told "query length <= 16" rather than that the package is missing. The routing was right and the diagnostic was actively misleading. Both gates now arm only when AITER could have run. Pre-existing for the flat-gather gate since #362; fixed for both rather than replicated. test_batch_decode_aiter_vs_fa2 parametrised (8,8), (16,4), (32,8) -- GQA ratios 1, 4, 4. Ratio 8 had no numerical coverage anywhere, and it is exactly where the decode sweep measured AITER's largest advantage, so "7.7x faster" rested on a kernel nothing had checked for correctness. Added (64,8): 144 cases pass, so the figure stands on a verified kernel. This is the one finding that could have invalidated a claim in the PR body rather than just tidying one. The two documented benchmark passes shared the default "decode" label, and the timing CSV is opened "w" -- the second pass overwrote the first. Distinct labels, and the reason recorded next to them. --kv-lens help advertised the pre-narrowing default; "1,,8" parsed as [1, 8] rather than erroring, silently sweeping a grid nobody asked for. Moved the both-gates-armed check to the top of the selector: it is argument validation and should not depend on which constraint the elif chain happens to answer first. Co-Authored-By: Claude <noreply@anthropic.com>

…ing question by measuring it (#373) ## Summary Takes the gfx942 ragged short-query win that #372 measured and deliberately left unclaimed, and settles the decode routing question #362 left open — by measuring it, and concluding that no decode gate should be built. ## What changed ### Routing - **`flashinfer/rocm/arch_caps.py`** — new `_AITER_RAGGED_GATED_Q_LEN` / `aiter_ragged_gated_q_len`, `{"gfx942": 16, "gfx950": None}`. Separate from the flat-gather table because the mechanism is different and gfx950's answer is "never". - **`flashinfer/rocm/prefill.py`** — `_aiter_ragged_short_query` / `_ragged_short_query_reason`, a `ragged_q_len=` kwarg on `_auto_select_prefill_backend` evaluated last in the elif chain, and the ragged wrapper's own `_backend_short_query_demoted` lifecycle. ### Benchmark - **`benchmarks/rocm/bench_batch_decode.py`** — `_KV_LENS` extended down to 128, a query-head axis, and `--batches` / `--kv-lens` / `--qo-heads` / `--output-dir`. ### Docs - **`docs/rocm/backends.md`** — the ragged gate, and the decode exemption restated as a measurement rather than an assertion. ## Architecture / design notes Arming the ragged path was not a matter of passing `max_q_len=`. Three things had to be their own: | | why | | :--- | :--- | | table | gfx950's row is `None` — a measured verdict, not a gap. Both accessors are bare `dict.get`, so an explicit `None` is indistinguishable from a deleted row; the test asserts **membership**. | | reason string | the flat-gather reason blames a copy of the whole KV cache, which ragged never makes. The paged demotion site matches its own string by equality, so sharing one would misreport *and* cross-trigger. | | demotion state | every `_backend_short_query_demoted` site was in the paged wrapper. Without a reset, one short extend strands a wrapper on `fa2` for every later long prefill. | The gate sits **last** in the selector's elif chain, where #362 put its sibling: ahead of it, a short batch would never reach the capability row or the dtype/layout checks, and `backend_fallback_reason` would report "too short" on a machine with no AITER installed. **A wrapper-scoped gate fires at most once, and that is the whole bug worth reading this PR for.** The first implementation put the gate inside `if self._backend == "auto"`, so a wrapper that resolved to AITER on a long prefill never re-entered the block — every later short plan went to AITER ungated. That is the long-then-short order, which is precisely what this gate exists for: a chat turn on a cached prefix *is* a short query after a long one, and it is the 4.74× cell. Short-then-long worked, and all four original tests happened to plan short first, so it shipped green through a full suite. The paged wrapper avoids this only because its native-paging probe gives it a second re-check; ragged has no probe, so the re-check is explicit. Cudagraph is asymmetric and deliberately so: demotion fires freely on a first plan, re-promotion never does, and the late re-check stays off under capture. Re-promoting swaps `_cached_module`/`_plan_info` that a captured graph still points at. Blocking demotion too — the obvious reading — would make the gate inert wherever graphs are used. Under capture the first plan decides, now documented alongside the paged wrapper's mirror-image exception. ## Benchmark results ### Ragged prefill, gfx942 (from #372, the basis for the threshold) All five measured shapes lose at 16 query tokens; first win at q24. | shape | AITER/fa2 at q16 | | :--- | :--- | | bs1 / kv1024 | 1.18× | | bs1 / kv8192 | 4.74× | | (five shapes, range) | 1.18–4.74× | gfx950 has a shape (bs32/kv2048) favouring AITER at *every* query length measured, so no gfx950 threshold serves and none is set. ### Batch decode, fa2 vs AITER — the measurement that says *not* to gate Swept at `0c43abef` on MI300X and MI350X, AITER `0.1.20+rocm10.1.0a20260819.3135022` on both, bf16 / head_dim 128 / 8 KV heads / page_size 16, medians over a 1 s repeat budget. | arch | qo heads | cells | AITER loses in | worst for AITER | best for AITER | | :--- | :--- | :--- | :--- | :--- | :--- | | gfx942 | 32 | 30 | 16 | 1.70× (b128/kv128) | 1.6× faster | | gfx950 | 32 | 30 | 15 | 1.16× (b256/kv128) | 1.7× faster | | gfx942 | 64 | 30 | **0** | 0.65× | **6.5× faster** | | gfx950 | 64 | 30 | **0** | 0.86× | **7.7× faster** | Three results, none of which supports a gate: - **There is no AITER decode "floor" at 0.16 ms.** AITER's fixed per-call cost is ~49–51 µs against fa2's ~38–44 µs. The entire losing region is that ~10 µs gap showing up while both kernels are launch-bound. - **The losing region is bounded by head count, not length.** At 64 query heads AITER wins every cell on both architectures. A gate keyed on `batch × kv_len` — which the bandwidth model predicts and which an earlier revision of this work proposed — would mis-route 70B and 405B decode by up to 7.7×. - **What is left to win is ~5–26 µs per call**, on h=32 only. Every launch swept was eager, and eager is also where `auto` mostly reaches AITER: under graph capture it resolves to `fa2` unless the caller passes `max_seq_len`, which vLLM and SGLang do not. Graph-mode and explicit-`aiter` launches were not timed. Against a 7× downside, that is not a trade worth making. Recorded rather than acted on: **fa2's decode kernel is capped by query head count and AITER's is not** — fa2 holds ~5.3 TFLOPS at h=64 on both arches regardless of shape while AITER reaches 37–41 TFLOPS, e.g. 6.41 ms vs 0.92 ms at b256/kv4096 on gfx950. KV traffic is identical at h=32 and h=64, so this is not bandwidth. Root cause unprofiled; it is a separate item, not this PR's. Those 64-head timings had no correctness coverage behind them when first posted: `test_batch_decode_aiter_vs_fa2` parametrised GQA ratios 1, 4 and 4, so ratio 8 — where AITER's advantage is largest — was unverified, and a fast-but-wrong kernel would have read as a win. `(64, 8)` is now in the matrix and its 144 cases pass, so the figures above stand on a checked kernel. ## Test plan - [x] `test_arch_caps.py::TestAiterRaggedQLenGate` — per-arch values, unknown-arch disarm, and gfx950 **membership** so deleting the row cannot pass - [x] `test_prefill_decode_dispatch.py` — selector arms on gfx942, falls through on gfx950, and the two gates' reason strings differ - [x] `test_batch_prefill_kernels.py` — ragged declines, **long-then-short still demotes**, demotion is not sticky, no re-promotion under cudagraph, explicit `backend="aiter"` survives, and gfx950 keeps AITER - [x] A/B three ways: disarming the table fails `test_gfx942_gates_through_16`; removing `ragged_q_len=self._max_q_len` fails 3 of 4 wrapper tests; removing the late re-check fails exactly `test_ragged_long_then_short_still_demotes` and nothing else - [x] Blast-radius suite (prefill, decode, dispatch, auto-fallback, arch-caps) — caught `test_ragged_prefill_auto_demotes_to_fa2` planning at exactly the new threshold and asserting on the wrong reason; fixed the same way #362 fixed its paged sibling - [x] `test_batch_decode_aiter.py` — GQA ratio 8 added to the parity matrix (144 cases), closing the gap under this PR's own 64-head measurements - [x] gfx942 (MI300X) and gfx950 (MI350X) - [x] `pre-commit run -a`, `/code-review xhigh` Benchmark-script defaults are left as they were: every config is materialised before the first run, so a bare invocation is ~55 GiB resident and extending `kv_lens` to 16384 would take it to ~116 GiB. The short-KV/multi-head sweep runs behind the flags in two passes, ~23 and ~45 GiB. ## Rollout Both behaviour changes are silent to callers: a `backend="auto"` ragged wrapper at `max_q_len <= 16` on gfx942 now gets `fa2`, with no API change and no opt-out short of `backend="aiter"`. `backend_fallback_reason` becomes non-`None` where it was `None` for those calls. --------- Co-authored-by: Claude <noreply@anthropic.com>

Copilot AI balanced review requested due to automatic review settings September 12, 2026 03:33

Copilot started reviewing on behalf of demandal25 September 12, 2026 03:34 View session

Copilot AI reviewed Sep 12, 2026

View reviewed changes

Comment thread flashinfer/rocm/decode.py Outdated

Comment thread flashinfer/rocm/decode.py Outdated

Comment thread flashinfer/rocm/prefill.py

Comment thread flashinfer/rocm/prefill.py Outdated

demandal25 and others added 19 commits September 12, 2026 14:28

Copilot AI review requested due to automatic review settings September 13, 2026 04:37

demandal25 force-pushed the rocm-spec-decode-verify branch from 1742249 to dd4f512 Compare September 13, 2026 04:37

Copilot started reviewing on behalf of demandal25 September 13, 2026 04:39 View session