Route short-query paged prefill to fa2, at 1.19-6.08x over AITER - #367
Open
demandal25 wants to merge 10 commits into
Open
Route short-query paged prefill to fa2, at 1.19-6.08x over AITER#367demandal25 wants to merge 10 commits into
demandal25 wants to merge 10 commits into
Conversation
The C++ half of fp8 paged prefill. No Python caller yet, so behaviour is
unchanged; bf16 and fp16 keep taking the same .so and the same qscale_type=0.
Everything here was read off the installed amd-aiter 0.1.20 tree rather than a
source checkout, per the ABI rule in CLAUDE.md:
- The fp8 variant exists and its .so is named
`mha_batch_prefill_fp8bf16_..._ndropout_pertensor_nsink.so` -- the dtype
segment spells both halves because AITER has no fp8-output prefill kernel,
and the qscale segment flips from `nqscale` to `pertensor`.
- `quant_scale_enum::pertensor = 1`, from
`aiter_meta/3rdparty/composable_kernel/example/ck_tile/01_fmha/quant.hpp`.
- fp8 *requires* descales: the `nqscale` fp8 variant compiles and then resolves
to no kernel ("no matching kernel found ... dtype=fp8bf16").
- `mha_batch_prefill_args` already declared q/k/v_descale_ptr; they were pinned
to nullptr.
The qscale token sits mid-name rather than in the trailing suffix, so
`build_so_name` gains an infix parameter and emits the token itself. That is why
the three call sites move their suffix literals.
**Per-tensor descales are enforced at the boundary, not coerced.** A per-head
descale of shape [8] is silently *accepted* by the pertensor kernel, which reads
element 0 -- so passing one applies head 0's scale to every head and returns
plausible wrong numbers. `TORCH_CHECK` rejects anything that is not a single
float32. This matters because `prefill.py` currently defaults `scale_q` to a
per-head `torch.ones(q.shape[1])`.
Verified: a bf16 paged prefill through `BatchPrefillWithPagedKVCacheWrapper`
still builds and runs after the change (gfx942, ROCm 10.0, aiter 0.1.20), so
the shared .so-name and enum paths are unregressed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…on fa2
Second half of the fp8 paged-prefill path. Not yet reachable end to end -- see
the blocker below, which this commit establishes rather than fixes.
Plumbing:
- `_auto_select_prefill_backend` gains `allow_fp8`, set only by the paged
wrapper. Single and ragged still take mha_fwd/mha_varlen_fwd, which have no
fp8 wiring, so admitting fp8 there would route to a kernel that cannot serve.
- `plan()` derives the output dtype from the query dtype: bf16 for fp8, since
AITER has no fp8-output prefill kernel. The module is now requested with that
dtype instead of the query dtype -- a no-op for fp16/bf16.
- `run()` takes keyword-only `scale_q`/`scale_k`/`scale_v` descales and hands
them to the shim. Keyword-only and appended, which `scripts/rocm_api_parity.py`
permits for a ROCm-only parameter ("fine appended or keyword-only").
- The AITER bootstrap probe passes descales for an fp8 dtype. Without them it
proves the wrong thing: the .so builds and dispatch then finds no kernel.
- The config .inc template includes the HIP fp16/bf16/fp8 headers. It is
included first in the generated .cu, so an fp8 config previously failed with
`unknown type name '__hip_fp8_e4m3_fnuz'` before anything else was parsed.
`_reject_fp8_on_fa2` closes the ugly half of the current behaviour. fa2 has no
fp8 kernel -- `include/flashinfer/rocm/attention/prefill.cuh:108` rejects 8-bit
types with a static_assert -- so an fp8 prefill used to surface as a ninja log:
RuntimeError: Ninja build failed. Ninja output:
prefill.cuh:108:17: error: static assertion failed due to requirement
'sizeof(__hip_fp8_e4m3_fnuz) != 1': 8-bit types not supported for CDNA3
It is now a NotImplementedError naming fp8 and pointing at the paged wrapper.
**Blocker, measured on gfx942 / aiter 0.1.20.** fp8 and bf16 have *different*
native paged page-size sets, and `_aiter_native_page_sizes()` returns one
hardcoded set for both:
bf16 native at {1, 16, 1024} (measured in the P0 work)
fp8 native at 16; 1024 is refused:
"no matching kernel found. page_size=1024, num_pages=1, dtype=fp8bf16"
code claims {128, 256, 1024} for every dtype
So every fp8 page size either falls outside the claimed set (16 -> flat gather,
which routes to mha_varlen_fwd and has no fp8 wiring) or is inside it and fails
the probe (1024). Either way the wrapper demotes to fa2 and hits the
static_assert above. The fp8 .so itself is fine -- it builds, loads, and appears
in the failing call's own stack trace.
Making that set correct *and dtype-aware* is a prerequisite for this feature,
not the separate no-op cleanup it was previously filed as.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The blocker was not page size and not a dtype-specific kernel set, as the
previous commit guessed. It is that **fp8 has no LSE instance**:
dtype page return_lse -> result
bf16 16 True -> ok fp8 16 True -> no matching kernel
bf16 1024 True -> ok fp8 1024 True -> no matching kernel
bf16 16 False -> ok fp8 16 False -> ok
`_aiter_native_paging_available` bootstraps both LSE variants and treats any
failure as "this config has no native paging", so fp8 was demoted on every page
size, fell through to flat-gather (mha_varlen_fwd, no fp8 kernel), and finally
to fa2's static_assert. The probe now asks only for the variants the dtype can
serve, and `run()` rejects fp8 + return_lse outright rather than degrading.
Sweeping the real kernel per dtype also corrects the capability set: it is
{1, 16, 1024} for bf16, fp16 **and** fp8 alike -- not the {128, 256, 1024} the
code claimed, which named the one size that works for none of them.
Capability and routing are now separate, because they disagree. P0.1 measured
the flat gather equal to or faster than native paging for bf16 at every batch
size, so `_aiter_paged_route_page_sizes` keeps fp16/bf16 on exactly the route
they take today and admits fp8 to native, which is the only route it has.
Widening it for fp16/bf16 is a benchmark, not a one-line edit.
Measured on gfx942 / MI300X, through BatchPrefillWithPagedKVCacheWrapper --
not the AITER-direct ceiling -- page_size=16, GQA 32/8, causal, quantisation
outside the timed region:
s_qo s_kv bs | bf16 ms fp8 ms | speedup
512 512 1 | 0.100 0.060 | 1.66x
1024 1024 1 | 0.126 0.082 | 1.53x
1024 1024 8 | 0.332 0.272 | 1.22x
2048 2048 4 | 0.457 0.354 | 1.29x
4096 4096 2 | 0.734 0.566 | 1.30x
Both dtypes resolve to `aiter`; the gap is narrower than the 1.45-1.70x
kernel-level ceiling because bf16 keeps the gather while fp8 goes native.
Correctness against an fp32 reference, same wrapper: max abs err 0.17-0.22 for
fp8 versus 0.02 for bf16, the expected magnitude for uncalibrated per-tensor
descales on random data.
Co-Authored-By: Claude <noreply@anthropic.com>
… at routing Nine tests for the new path, including the ones that keep it honest: - `test_fp8_ignoring_descales_would_be_caught` is the A/B for the numeric test. AITER accepts a descale it does not honour elementwise, so "fp8 matches the reference" only means something if a 4x q_descale actually moves the result. - `test_fp8_rejects_per_head_descale` pins the boundary check. The per-tensor kernel reads element 0 of whatever it is handed, so a per-head tensor would apply head 0's scale to every head and return plausible wrong numbers. - `test_fp8_rejects_return_lse` and `test_single_prefill_fp8_raises_instead_of_a_ninja_log` pin the two error paths, the second because fa2's refusal is a static_assert and used to reach the user as a compiler log. - The numeric test calibrates against the bf16 error on the same inputs rather than a constant: fp8 error is dominated by the uncalibrated per-tensor descale, so a fixed bound would either pass anything or fail on noise. Three existing tests in test_batch_prefill_kernels.py selected page sizes from `_aiter_native_page_sizes()`, which was fine while capability and routing were the same set. They now use `_aiter_paged_route_page_sizes`, since what they are actually asserting is which route a call takes: - `test_paged_softcap_guard_tracks_the_paging_route` **failed** before this change -- it planned at page size 1, which is capable but not routed for fp16, so the gather guard fired where the test expected silence. - Its kv_len moves 512 -> 1024 and the fallback test's page size 128/256 -> 1024, because otherwise both merely *skip* under the corrected set. A test that stops running is worse than one that fails. gfx942: 13 passed, 0 skipped across the soft-cap, native-paging and strict-mode selection; 9 passed for the new file. gfx950: 9 passed for the new file. Co-Authored-By: Claude <noreply@anthropic.com>
Committed unformatted: the pre-commit output that would have caught it was swallowed by a background launch in the same command. Co-Authored-By: Claude <noreply@anthropic.com>
…g encoding
`/code-review xhigh` returned 15 findings. Six were real defects in paths this
PR added, all of the shape it set out to remove -- a silent wrong answer or a
compiler log where an error belongs.
**Wrong fp8 encoding was accepted and returned NaN.** Both e4m3fnuz and e4m3fn
are 8 bits and neither AITER nor the .so name distinguishes them, so the
non-native one is read under the wrong exponent bias. Measured on gfx942:
e4m3fnuz max|diff vs native| = 0.00000
e4m3fn max|diff vs native| = nan
`_require_native_fp8_dtype` now rejects it, taking the arch's encoding from
`aiter.dtypes.fp8` rather than duplicating the table -- AITER picks it per
architecture and its kernels are compiled against that choice.
**The post-probe demotion bypassed the fa2 guard.** `plan()` checks fp8-on-fa2
before the native-paging probe, then demotes to fa2 forty lines later if the
probe fails; an fp8 call at a non-routed page size went straight to the
`static_assert` and printed ninja output. Re-guarded at both demotion sites,
paged and ragged.
**The flat-gather branch accepted fp8 and dropped the descales** it had just
validated -- it dispatches mha_varlen_fwd, which has no fp8 kernel. Now a
TORCH_CHECK.
**`partial_state` defeated the LSE guard** by allocating `lse` itself after the
check, asking for an fp8 LSE variant plan() never bootstraps and AITER never
builds. Folded into the same guard.
**Descales were not device-checked**, so a CPU scalar's host pointer reached
the kernel. **`_cached_o_data_type`** raised AttributeError on run()-before-
plan() where upstream uses getattr. **The dlopen hints** printed "bf16" for an
fp8 variant, advising a rebuild of the wrong kernel.
Also: `scale_q/k/v` are rejected rather than silently dropped when a call
resolves to fa2; `test_prefill_decode_dispatch.py` still asserted the old
{128, 256, 1024} set (4 failures); `bench_aiter_prefill.py` picked its "native"
page size from capability, which after the correction resolved to 16 and would
have labelled two flat-gather rows as native; two new tests were tautological
or missing their aiter skip.
gfx942: 44 passed across tests/rocm/test_fp8_paged_prefill.py and
tests/rocm/test_prefill_decode_dispatch.py. pre-commit clean.
Co-Authored-By: Claude <noreply@anthropic.com>
…plan
AITER's mha_batch_prefill costs the same whatever the query length -- its
median time is flat from 16 to 256 query tokens at a fixed batch and kv --
so a batch of a few query rows pays a full KV scan for them. `auto` picked
it anyway, at every shape, with an empty backend_fallback_reason.
Measured through BatchPrefillWithPagedKVCacheWrapper, page_size 16, GQA
32/8, hd128, causal bf16, aiter/fa2 median ratio (>1 = AITER slower):
gfx942 q16 q24 q32 q64 q128 q256
bs1/kv1024 1.88 1.84 1.83 1.61 1.43 0.85
bs1/kv8192 5.04 3.83 3.79 1.87 1.03 0.55
bs8/kv1024 1.46 0.96 0.88 0.49 0.36 0.27
bs8/kv8192 2.10 1.09 1.08 0.54 0.32 0.22
bs32/kv2048 1.86 1.04 1.02 0.54 0.30 0.25
gfx950 q16 q24 q32 q64 q128 q256
bs1/kv1024 1.62 1.71 1.51 1.42 1.17 0.74
bs1/kv8192 6.08 4.67 3.64 1.83 1.00 0.52
bs8/kv1024 1.50 1.00 0.88 0.47 0.28 0.22
bs8/kv8192 2.06 1.04 0.89 0.48 0.24 0.15
bs32/kv2048 1.19 0.70 0.67 0.38 0.26 0.21
16 is the last column where AITER loses in every cell on both
architectures; at 24 it already wins at bs8/kv1024 on gfx942 and
bs32/kv2048 on gfx950. The threshold is deliberately the conservative end
of a wide crossover (bs1 does not cross until ~192) because the error is
asymmetric: choosing fa2 wrongly costs up to 4.6x at bs8/kv8192/q256,
choosing AITER wrongly costs at most the 1.9x bs1 cells. 16 also covers
the case that motivated this -- speculative-decode verify, where q is the
draft length plus one.
Ragged prefill is deliberately not gated. The same sweep through
mha_varlen_fwd splits by architecture at q16: gfx942 loses 1.07-4.42x but
gfx950 runs bs32/kv2048 at 0.59, i.e. 1.7x *faster* than fa2. A shared
threshold would regress CDNA4 there.
The gate needs the paged wrapper to stop caching its first answer. plan()
resolved `auto` once and kept it, so a served wrapper's backend was decided
by whichever shape it happened to see first -- which for a shape-dependent
rule is worse than no rule. `auto` now re-resolves every plan, except under
CUDA-graph capture, where the captured graph holds buffers belonging to the
backend it captured.
Speed is also the first reason that is not a capability, so it is last in
the chain (a constraint the caller controls should out-rank it in the
message) and logs at info rather than warning -- picking the faster kernel
is not a degradation, and a warning would fire on every server start.
Raw CSVs: tmp/bench/xover_{gfx942,gfx950}.csv and xoverR_* for ragged.
Co-Authored-By: Claude <noreply@anthropic.com>
Unit coverage of the selector goes in the dispatch file, which launches no kernel: the threshold boundary in both directions, an unspecified qo_len staying ungated (single prefill passes none), a capability reason out-ranking the speed one, and the absence of a warning. The wrapper-level file is the one that would have caught the sticky backend. test_planning_longer_returns_to_aiter fails on the old code -- the first plan's fa2 persisted -- and so does its mirror, since the old code never re-entered the selector at all. The graph-mode case plans long first, because graph mode caps later plans at the first plan's row count. The AITER-available skip is autouse and checked before every assertion: without AITER every plan resolves to fa2 and the whole file would pass for the wrong reason. Co-Authored-By: Claude <noreply@anthropic.com>
`/code-review xhigh` on 864b4a9..d96f431. The six that mattered: **fp8 lost its only route.** fa2 has no fp8 prefill kernel, so preferring it on speed turned a working short-query fp8 plan into NotImplementedError with nothing left to fall back to. Reproduced before the fix: RAISED: NotImplementedError fp8 prefill (dtype=torch.float8_e4m3fnuz) has no in-tree fa2 kernel on ROCm ... after: PLAN OK, backend = aiter **Two existing tests broke and had not been run.** Both plan a paged `auto` wrapper at exactly 16 query tokens and assert AITER: test_batch_prefill_auto_selects_aiter (12 parametrizations) and test_paged_prefill_auto_demotes_to_fa2, whose subject -- the bootstrap-failure demotion -- was no longer reached at all. Their qo_len moves to 32, the smallest value clearing the gate that keeps qo_len < kv_len. **The gate out-ranked the importability check.** It returned before `_aiter_ops_importable()`, so a box with no amd-aiter, or one below the ABI floor, reported a perf choice at info level and never logged the warning that says AITER is missing. Moved after both, which is what the comment claimed. **Two logs became unbounded.** The reason string embeds the length, so keying the warn-once set on it gave one entry and one line per distinct query length; the speed decline now dedups on a constant. Separately, per-plan re-resolution re-enters the soft-cap flat-gather demotion every plan, whose warning had no dedup because the old sticky backend meant it could only fire once -- now routed through _aiter_auto_warned like every other site. Tests gained the missing controls: the wrapper file's AITER check is now a positive control that plans a long query and skips if the probe demotes, rather than trusting capability; TestNumerics asserts the two wrappers really landed on different backends before comparing them; the dispatch tests skip on a gated arch instead of failing; and the graph-mode case asserts the selector did not run at all rather than that it happened to re-pick AITER. Also declined from that review: moving the ragged decision into arch_caps. There is no arch-differing behaviour to express -- ragged is ungated on both -- and adding a per-arch threshold is a new feature needing its own sweep. Co-Authored-By: Claude <noreply@anthropic.com>
… tokens Found by the pinned regression run at 678ad84, on both architectures, after the review caught the other two. Same defect: _plan_softcap_flat_gather builds a qo_len=16 batch, so the `auto` arm now declines on speed before reaching the soft-cap demotion it asserts on. AssertionError: assert 'logits_soft_cap' in 'qo_len=16 <= 16 (fa2 is faster than AITER for short queries)' qo_len goes to 32; it is not part of the defect shape (causal, cap>0, head_dim=128, kv_len>=512), so the test's subject is unchanged. The `aiter` arm never regressed -- an explicit backend skips the gate. The three other paged tests at qo_len=16 are unaffected because they request backend="aiter" outright: test_batch_prefill_aiter_flat_gather_bf16 and the two bootstrap-rejection tests. Checked rather than assumed. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
The routing gate extends beyond the measured kernel path, and several FP8 edge paths do not enforce the documented failure behavior.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Routes short paged-prefill queries to FA2, adds FP8 AITER paged-prefill support, and updates routing tests, benchmarks, and documentation.
Changes:
- Re-selects paged-prefill backends per plan with a 16-token FA2 threshold.
- Adds FP8 inputs, bf16 outputs, per-tensor descales, and AITER dispatch support.
- Expands routing, numerical, fallback, and benchmark coverage.
File summaries
| File | Description |
|---|---|
flashinfer/rocm/prefill.py |
Implements routing and FP8 handling. |
flashinfer/rocm/arch_caps.py |
Updates backend capability notes. |
include/flashinfer/rocm/attention/aiter/batch_prefill.cuh |
Passes FP8 descales to AITER. |
include/flashinfer/rocm/attention/aiter/aiter_loader.h |
Extends AITER variant keys. |
csrc/rocm/batch_prefill_paged_aiter.cu |
Adds FP8 paged dispatch and validation. |
csrc/rocm/batch_prefill_paged_aiter_jit_pybind.cu |
Extends the native binding signature. |
csrc/rocm/batch_prefill_aiter_customize_config.jinja |
Adds HIP FP8 headers. |
csrc/rocm/aiter_loader.cc |
Resolves FP8/scaled AITER libraries. |
tests/rocm/test_short_query_prefill_routing.py |
Tests routing and replanning. |
tests/rocm/test_prefill_decode_dispatch.py |
Tests selector boundaries and precedence. |
tests/rocm/test_fp8_paged_prefill.py |
Tests FP8 routing and numerics. |
tests/rocm/test_batch_prefill_kernels.py |
Updates paging-route tests. |
tests/rocm/test_aiter_auto_fallback.py |
Adjusts fallback test shapes. |
benchmarks/rocm/testlist_rocm.txt |
Adds short-query benchmark cases. |
benchmarks/rocm/bench_aiter_prefill.py |
Uses routed page-size selection. |
docs/rocm/backends.md |
Documents routing and FP8 constraints. |
README.md |
Updates the generated support matrix. |
Review details
Suppressed comments (3)
flashinfer/rocm/prefill.py:519
- This guard recognizes only the two E4M3 encodings. Other PyTorch FP8 dtypes such as
torch.float8_e5m2are rejected by the selector, then still reach fa2 because this check misses them; the stated compiler-log protection is therefore bypassed. Detect every float8 dtype before dispatching to fa2.
if backend == "fa2" and dtype_q in FP8_PREFILL_DTYPES:
flashinfer/rocm/prefill.py:1822
- An explicit
backend="aiter"bypasses this fa2-only guard, so fp8 single prefill proceeds into a shim that accepts only fp16/bf16 instead of raising the promised route-levelNotImplementedError. Reject fp8 for the single-prefill AITER route before bootstrapping it.
_reject_fp8_on_fa2(q.dtype, backend)
flashinfer/rocm/prefill.py:3822
- An explicit ragged
backend="aiter"bypasses this fa2-only guard and reaches the fp16/bf16-only ragged shim. This contradicts the documented guarantee that non-paged fp8 routes raise a clearNotImplementedError; reject fp8 before selecting either ragged backend.
_reject_fp8_on_fa2(q_data_type, self._backend)
- Files reviewed: 17/17 changed files
- Comments generated: 5
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+504
to
+510
| native = _native_fp8_dtype() | ||
| if dtype_q in FP8_PREFILL_DTYPES and native is not None and dtype_q != native: | ||
| raise NotImplementedError( | ||
| f"fp8 prefill needs this GPU's encoding, {native}; got {dtype_q}, " | ||
| "which the kernel reads under the wrong exponent bias and returns " | ||
| "NaN for. Re-quantize with aiter.dtypes.fp8." | ||
| ) |
| causal=causal, | ||
| logits_soft_cap=logits_soft_cap, | ||
| kv_len=softcap_kv_len, | ||
| qo_len=self._max_q_len, |
Comment on lines
+2989
to
+2991
| scale_q: Optional[torch.Tensor] = None, | ||
| scale_k: Optional[torch.Tensor] = None, | ||
| scale_v: Optional[torch.Tensor] = None, |
| global PAGE | ||
| original, PAGE = PAGE, 32 | ||
| try: | ||
| with pytest.raises((NotImplementedError, RuntimeError, ValueError)): |
Comment on lines
+533
to
+536
| * **Descales are required and must be per-tensor** — a single float32 each for | ||
| `scale_q`, `scale_k`, `scale_v`, passed to `run()`. AITER reads element 0 of | ||
| whatever it is given, so a per-head tensor would silently apply head 0's scale | ||
| to every head; the shim rejects it instead. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
backend="auto"sent every paged prefill to AITER regardless of query length, and at short queries AITER is 1.19-6.08x slower than the in-treefa2kernel on both architectures — with an emptybackend_fallback_reason, so it read as a deliberate choice rather than the oversight it was. This gates the paged route at 16 query tokens, which is where speculative-decode verify and the tail of a chunked prefill live, and makesautore-decide on everyplan()instead of caching whatever the first plan saw.Stacked on #366 — the diff shows that PR's commits until it merges. The four commits here are the last four.
What changed
flashinfer/rocm/prefill.py—_auto_select_prefill_backendtakesqo_lenand declines AITER at or below_AITER_SHORT_QO_LEN(16), after the capability and importability checks and never for fp8. The pagedplan()passesself._max_q_len(already computed for the kernel) and re-resolvesautoevery call. Ragged passes nothing, deliberately.tests/rocm/test_batch_prefill_kernels.py,tests/rocm/test_aiter_auto_fallback.py— three existing tests planned anautopaged wrapper at exactly 16 query tokens and asserted a backend or a reason; theirqo_lenmoves to 32, the smallest value clearing the gate that keepsqo_len < kv_len. Three other paged tests at 16 are unaffected because they requestbackend="aiter"outright, which skips the gate — checked, not assumed.benchmarks/rocm/testlist_rocm.txt— three short-query rows plus aq=64control that must stay on AITER. These are atpage_size=16, the serving default, unlike the existing paged rows at 128.docs/rocm/backends.md— a "short queries go tofa2" subsection under batch prefill, with the measured ratios and the reason ragged is excluded.tests/rocm/test_short_query_prefill_routing.py(new) — wrapper-level routing, re-planning in both directions, the graph-capture exception, ragged staying on AITER, and numerics.tests/rocm/test_prefill_decode_dispatch.py— selector-level unit coverage: the boundary in both directions,qo_len=Nonestaying ungated, precedence against a capability reason, and no warning.Architecture / design notes
AITER's cost does not depend on query length. At a fixed batch and kv,
mha_batch_prefill's median time is flat from 16 to 256 query tokens — 0.079/0.079/0.080/0.079/0.079/0.078/0.078/0.078/0.077 ms at bs1/kv1024 on gfx942.fa2's grows with the query. The ratio therefore crosses 1 wherever fa2's growing cost meets AITER's fixed floor, which is why this is a threshold onqo_lenand not on total tokens or arithmetic intensity.16 is the conservative end of a wide crossover, chosen for the asymmetry. The crossover is not a single value: bs≥8 crosses between 24 and 32, but bs1 does not cross until ~192 because fa2 is latency-bound there and AITER is occupancy-bound. No one-parameter law fits all 90 points — total query tokens is refuted by bs32/q16 (512 tokens, AITER loses) against bs8/q24 (192 tokens, AITER wins). So the threshold is placed where every measured cell agrees, and 16 is the last such column. That is the right side to err on: choosing fa2 wrongly costs up to 4.6x (bs8/kv8192/q256), choosing AITER wrongly costs at most the 1.9x bs1 cells. It also happens to be exactly the spec-decode window, where
qo_lenis the draft length plus one.Ragged is deliberately not gated, because the same sweep splits by architecture. Through
mha_varlen_fwdat 16 query tokens, gfx942 loses 1.07-4.42x but gfx950 runs bs32/kv2048 at 0.59 — 1.7x faster than fa2, driven by an AITER kernel that is 2.3x quicker there than CDNA3's (0.185 ms vs 0.421 ms). A shared threshold would regress CDNA4. Extending the gate to ragged needs a per-arch table, and that needs more than one contradicting cell to justify.The gate is useless without per-plan re-resolution.
plan()resolvedautoonce and kept the answer for the life of the wrapper, so a served wrapper's backend was decided by whichever shape it saw first. For a shape-dependent rule that is worse than no rule: a wrapper that started on a decode-ish step would stay on fa2 through every later chunked prefill, at up to 4.6x.autonow re-resolves each plan. Under CUDA-graph capture it does not, because the captured graph holds buffers belonging to the backend it captured — that case keeps the old sticky behaviour and is covered by a test.Speed is the first reason in the chain that is not a capability, which decides three things about where it sits. It runs after the capability checks and after the
_aiter_ops_importable()check, so a user whose AITER is missing or below the ABI floor hears that rather than a claim that fa2 happens to be quicker. It logs atinforather thanwarning, and dedups on a constant rather than on the message, since the reason embeds the length and would otherwise emit one line per distinct query. And it is skipped entirely for fp8: paged AITER is fp8's only prefill route (#366), so preferring fa2 there would turn a working call into aNotImplementedErrorwith nothing to fall back to.Benchmark results
Through
BatchPrefillWithPagedKVCacheWrapper,page_size=16, GQA 32/8, hd128, causal bf16,--refcheckclean on all 90 configs per architecture. Ratio isaiter / fa2median time measured back-to-back in one process, so >1 means AITER is slower and the gate is a win of that factor. The boldq16column is the only one the gate touches; every column to its right is unchanged and shown to justify where the threshold sits.gfx942 / MI300X — paged
gfx950 / MI350X — paged
Every cell at q16 is above 1 on both architectures; at q24 two cells are already below it (gfx942 bs8/kv1024 at 0.96, gfx950 bs32/kv2048 at 0.70), which is what fixes the threshold at 16.
Ragged, at q16 only — the reason it is excluded
After the change, the same shapes through the testlist.
autolands on fa2 and tracks it within noise; the control keeps AITER and its 1.8-2.1x win.autobeforeautoafterRead the speedup from the per-run ratio tables above rather than across these two columns: fa2 at bs1/kv8192/q16 measured 0.068 ms in the sweep and 0.098 ms here, a run-to-run spread that the within-process ratio is immune to.
Those rows also now carry
backend_fallback_reason=qo_len=16 <= 16 (fa2 is faster than AITER for short queries), where before they were the empty-reason case the testlist header tells readers to investigate.Test plan
qo_len: 6 failures (4 routing, 2 re-planning)NotImplementedError; after,PLAN OK, backend = aitertest_batch_prefill_kernels.py,test_single_prefill_kernels.py,test_prefill_sliding_window.py,test_fp8_paged_prefill.py,test_aiter_auto_fallback.py,test_prefill_decode_dispatch.py,test_short_query_prefill_routing.py,test_batch_decode_aiter.pytest_single_prefill_kernels.py, whose large-head-dim failures there predate this branchpin-run.sh --checkconfirms neither tree moved during the run.test_benchmark_harness.py,test_rocm_small_surface.py,test_arch_caps.py: 55 passed--refcheckclean throughout/code-review xhigh— 15 findings; the six real defects are fixed in 678ad84 and described in its message, and the two it caught in existing tests were breakage this PR would otherwise have shippedpre-commit run -aKnown limitations
Single prefill passes no
qo_len: it reachesmha_fwd, a different kernel, and the threshold was not measured there. Ragged is excluded on evidence, above. Both would need their own sweep, and ragged additionally a per-arch table.