fp8 paged prefill on the AITER route, at 1.22-1.66x over bf16 by demandal25 · Pull Request #366 · AMD-Ecosystem/flashinfer

demandal25 · 2026-09-12T14:05:54Z

Summary

fp8 prefill was not merely ungated on ROCm — it was broken on every route, failing with a ninja compiler log rather than an error. This makes it work on the paged route, where AITER has a kernel, and makes every other route fail with a sentence instead of a build log. Measured 1.22–1.66× over bf16 on gfx942; on gfx950 the gain is large only at short sequences and flat by 4096, which is reported below rather than averaged away.

What changed

C++ shim

csrc/rocm/aiter_loader.cc, include/flashinfer/rocm/attention/aiter/aiter_loader.h — VariantKey gains kFp8Bf16 and a has_qscale axis. The qscale token sits mid-name, so build_so_name emits it rather than carrying it in a fixed suffix; the three call sites move their literals accordingly.
include/flashinfer/rocm/attention/aiter/batch_prefill.cuh — passes the descale pointers (previously pinned to nullptr) and selects quant_scale_enum::pertensor for fp8.
csrc/rocm/batch_prefill_paged_aiter.cu — accepts fp8 with per-tensor descales, enforces the bf16 output, and refuses fp8 on the flat-gather branch.
csrc/rocm/batch_prefill_aiter_customize_config.jinja — includes the HIP fp16/bf16/fp8 headers. The .inc names __hip_fp8_* types, so any fp8 config previously died with unknown type name before anything else parsed.

Python

flashinfer/rocm/prefill.py — allow_fp8 on the backend selector (paged only); bf16 output dtype for an fp8 query; keyword-only scale_q/scale_k/scale_v; an fp8-aware bootstrap probe; _reject_fp8_on_fa2 and _require_native_fp8_dtype; and _aiter_paged_route_page_sizes, which separates routing from capability.

Architecture / design notes

The blocker was that fp8 has no LSE instance, which took three wrong guesses to find. _aiter_native_paging_available bootstraps both LSE variants and treats any failure as "no native paging for this config", so fp8 was demoted on every page size, fell through to flat-gather (mha_varlen_fwd, no fp8 kernel), and landed on fa2's static_assert:

dtype	page	`return_lse`	result
bf16	16	True	ok
bf16	1024	True	ok
fp8	16	True	no matching kernel
fp8	1024	True	no matching kernel
fp8	16	False	ok

Capability and routing are now separate, because they disagree. Sweeping the real kernel per dtype, the native page-size set is {1, 16, 1024} for fp16, bf16 and fp8 alike — not the {128, 256, 1024} the code claimed, which names the one size that works for none of them. But being able to route natively is not a reason to: earlier measurement showed 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.

auto re-resolves on every plan(). It used to go concrete on the first call and never re-read _backend_requested, so a wrapper whose first plan hit any AITER constraint refused fp8 for the rest of its life — and a served wrapper is re-planned every step. Under CUDA-graph capture the first answer still sticks, because the captured graph holds that backend's buffers.

Per-tensor descales are enforced, not coerced. A per-head descale of shape [8] is silently accepted by the pertensor kernel, which reads element 0 — so it would apply head 0's scale to every head and return plausible wrong numbers. prefill.py still defaults scale_q to a per-head ones(q.shape[1]) elsewhere, which is exactly how that would have happened.

The wrong fp8 encoding returns 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 → 0.00000 against the native result, e4m3fn → nan. The arch's encoding comes from aiter.dtypes.fp8 rather than a table of our own, since AITER's kernels are compiled against that choice.

Everything above was read off the installed amd-aiter 0.1.20 tree, per the ABI rule in CLAUDE.md — including quant_scale_enum::pertensor = 1 and the mha_batch_prefill_fp8bf16_..._pertensor_nsink.so naming.

Benchmark results

Through BatchPrefillWithPagedKVCacheWrapper — not the AITER-direct ceiling — page_size 16, GQA 32/8, causal, bf16 vs fp8, quantisation outside the timed region. Both dtypes resolve to aiter.

gfx942 / MI300X

s_qo	s_kv	bs	bf16 ms	fp8 ms	speedup
512	512	1	0.100	0.060	1.66×
1024	1024	1	0.126	0.082	1.53×
1024	1024	8	0.332	0.272	1.22×
2048	2048	4	0.457	0.354	1.29×
4096	4096	2	0.734	0.566	1.30×

gfx950 / MI350X — the gain is front-loaded and flat by 4096

s_qo	s_kv	bs	bf16 ms	fp8 ms	speedup
512	512	1	0.137	0.076	1.80×
1024	1024	1	0.143	0.101	1.41×
1024	1024	8	0.249	0.237	1.05×
2048	2048	4	0.327	0.302	1.08×
4096	4096	2	0.488	0.481	1.01×

CDNA4's bf16 is already much faster at these shapes (0.488 ms vs CDNA3's 0.734 at 4096), so there is less headroom to recover. fp8 is never slower on either architecture, but on gfx950 it is worth having mainly for short sequences.

Accuracy, same wrapper against an fp32 reference: 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.

Test plan

gfx942 (MI300X) — test_fp8_paged_prefill.py, test_prefill_decode_dispatch.py, test_aiter_variants.py, test_aiter_auto_fallback.py, and the full test_batch_prefill_kernels.py: no F/E
gfx950 (MI350X) — the same five files
Numerics run at every routed fp8 page size (1, 16, 1024) against an fp32 reference, not just the serving default — page size is the AITER dispatch axis, so set membership is not coverage
A/B on the numeric test: a 4× q_descale must change the result, or "fp8 matches the reference" proves nothing — AITER accepts descales it does not honour elementwise
A/B on every guard added in review: dropping the call site must fail its test. Four for the route guards, two for the round-3 fixes
Both architectures benchmarked through the wrapper, table above
/code-review xhigh — 15 findings; the six real defects are fixed in 6dc0b0fc0
Copilot rounds 1–4: 13 inline comments and 5 suppressed, answered in 2ad5fa357 and 4ce38d89e; round 4 came back with 0 new
pre-commit run -a

Known limitations

fp8 is paged-native only. Single prefill, ragged prefill and the paged flat-gather path all reach mha_fwd/mha_varlen_fwd, which have no fp8 kernel; all three raise NotImplementedError naming fp8, regardless of backend — an explicit backend="aiter" used to reach AITER's own bootstrap and get RuntimeError: invalid argument for fmha_fwd. return_lse and partial_state are unavailable for fp8, and the output is always bf16.

An encoding that cannot be read from aiter.dtypes.fp8 is refused rather than guessed: the wrong encoding returns NaN, so failing open there is worse than failing.

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>

Copilot

🟡 Changes recommended

Explicit AITER requests can still reach unsupported FP8 single, ragged, or flat-gather bootstraps and expose raw failures.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds FP8 paged prefill through AITER with BF16 output, per-tensor descales, routing guards, tests, benchmarks, and documentation.

Changes:

Adds FP8-aware AITER dispatch and shared-library resolution.
Separates native paging capability from dtype-specific routing.
Adds correctness tests, benchmarks, and updated ROCm documentation.

File summaries

File	Description
`flashinfer/rocm/prefill.py`	Adds FP8 routing, validation, output dtype, and descales.
`flashinfer/rocm/arch_caps.py`	Updates AITER capability notes.
`csrc/rocm/aiter_loader.cc`	Resolves FP8 and qscale AITER variants.
`csrc/rocm/batch_prefill_paged_aiter.cu`	Adds FP8 paged execution and validation.
`csrc/rocm/batch_prefill_paged_aiter_jit_pybind.cu`	Extends the binding with descales.
`csrc/rocm/batch_prefill_aiter_customize_config.jinja`	Includes HIP FP8 types.
`include/flashinfer/rocm/attention/aiter/aiter_loader.h`	Extends variant keys for FP8 scaling.
`include/flashinfer/rocm/attention/aiter/batch_prefill.cuh`	Passes descales into AITER kernels.
`tests/rocm/test_fp8_paged_prefill.py`	Adds FP8 correctness and guard coverage.
`tests/rocm/test_batch_prefill_kernels.py`	Updates native-routing tests.
`tests/rocm/test_prefill_decode_dispatch.py`	Updates native page-size expectations.
`benchmarks/rocm/bench_aiter_prefill.py`	Updates benchmark route selection.
`docs/rocm/backends.md`	Documents FP8 constraints and paging behavior.
`README.md`	Updates the generated support matrix.

Review details

Suppressed comments (1)

flashinfer/rocm/prefill.py:3771

Explicit backend="aiter" is not rejected here, so fp8 ragged prefill proceeds to bootstrap the unsupported mha_varlen_fwd variant and may return a compiler/probe error instead of the documented NotImplementedError. Ragged fp8 should be rejected regardless of the selected backend.

            _reject_fp8_on_fa2(q_data_type, self._backend)

Files reviewed: 14/14 changed files
Comments generated: 6
Review effort level: Balanced

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

Copilot

🟡 Changes recommended

Explicit AITER paths can still expose unsupported FP8 build failures, and the variant-name test is incompatible with the updated loader signature.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (4)

flashinfer/rocm/prefill.py:2729

For explicit backend="aiter", an fp8 page size outside this route set—or a native probe that returns false—falls through to _aiter_bootstrap_batch_ragged_prefill because demotable is false. That unsupported fp8 varlen bootstrap can surface a Ninja build error before the new C++ flat-gather guard runs. Detect fp8 immediately when use_native_paging is false and raise the native-paging NotImplementedError before bootstrapping the flat-gather family.

                if page_size in _aiter_paged_route_page_sizes(q_data_type):

flashinfer/rocm/prefill.py:2988

The two run overloads above do not declare these new keyword arguments, and the Parameters section does not document them. Runtime calls work, but type checkers and generated API signatures reject run(..., scale_q=..., scale_k=..., scale_v=...), which is now the required fp8 interface. Add the three keyword-only tensor parameters to both overloads and document their per-tensor float32/device requirements.

        scale_q: Optional[torch.Tensor] = None,
        scale_k: Optional[torch.Tensor] = None,
        scale_v: Optional[torch.Tensor] = None,

flashinfer/rocm/prefill.py:1833

This guard only rejects fp8 after auto resolves to FA2 or when FA2 is explicitly selected. With backend="aiter", single prefill proceeds into _aiter_bootstrap_single_prefill_* even though AITER has no single-prefill fp8 kernel, so users still receive the compiler/build failure this PR intends to replace. Reject fp8 unconditionally for the single-prefill API before backend-specific bootstrapping.

    _reject_fp8_on_fa2(q.dtype, backend)

tests/rocm/test_fp8_paged_prefill.py:170

This test accepts any RuntimeError or ValueError, so the Ninja/compiler failure it is specifically meant to prevent would make the test pass. Assert the intended public exception and message instead, so a regression back to a build log is detectable.

        with pytest.raises((NotImplementedError, RuntimeError, ValueError)):

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

**The variant-store mirror broke on the merge, and its test fails outright.** #363 landed `tests/rocm/test_aiter_variants.py`, which parses the `build_so_name` call sites out of aiter_loader.cc and asserts they match the Python prebuild table. This PR gave `build_so_name` a fourth string argument, so the regex matched nothing: AssertionError: no build_so_name call sites parsed; the regex or the C++ moved `Family` now carries `infix` alongside `prefix`/`suffix`, `so_name()` places the quantisation token between them, and the regex takes four strings. All 40 generated names are byte-identical to before, so the store on disk stays valid. The table keeps `nqscale` hard-coded: `pertensor` is fp8, fp8 is batch-prefill only, and that family is already `servable_from_store=False`. **fp8 escaped its guards under an explicit `backend="aiter"`.** Both routes measured before the fix: ragged aiter+fp8 : RuntimeError: invalid argument for fmha_fwd paged aiter+fp8 page=128 : RuntimeError: invalid argument for fmha_fwd Copilot called this a ninja build error; it is not, but the remedy stands -- neither is the documented NotImplementedError. `_reject_fp8_off_native_paging` fires on both routes regardless of backend, before anything bootstraps the flat-gather family. After: ragged aiter+fp8 : NotImplementedError: ... ragged batch prefill paged aiter+fp8 page=128 : NotImplementedError: ... flat-gather route paged aiter+fp8 page=16 : NO RAISE **`scale_q`/`scale_k`/`scale_v` were invisible to type checkers.** They were on the paged `run` implementation but on neither `@overload`, and a checker resolves against the overloads only. Minimal repro, mypy 1.17.1: error: No overload variant of "run" of "W" matches argument types "int", "int" Both overloads now declare them, and the Parameters section documents the per-tensor float32 contract -- including that a per-head tensor is read as element 0 rather than rejected. Also: the non-routed-page-size test accepted `RuntimeError`/`ValueError`, which is exactly what a regression would raise, so it is now `NotImplementedError` with a message match; the benchmark queried bf16 routing while running fp16 and its header still advertised page 256 as native; and `docs/rocm/backends.md` said in three places that other page sizes "still work" and that the descales are dropped, both of which this PR makes false for the native paged route. Declined: rejecting fp8 unconditionally in single prefill. It already answers `RuntimeError: AITER backend supports fp16/bf16 only; got dtype=Float8_e4m3fnuz` on the aiter arm and NotImplementedError on the fa2 arm, so both are covered. Co-Authored-By: Claude <noreply@anthropic.com>

demandal25 · 2026-09-13T16:44:25Z

Suppressed comments from reviews 5186639352 and 5189588416, all in 2ad5fa3:

Item	Disposition
`prefill.py:3771` / `:2729` — explicit `aiter` reaches the fp8 flat-gather and ragged bootstraps	Fixed. Measured `RuntimeError: invalid argument for fmha_fwd`, not a ninja log, but the remedy stands
`prefill.py:2988` — `scale_*` missing from both `run` overloads	Fixed, plus the Parameters entry
`test_fp8_paged_prefill.py:170` — exception tuple too broad	Fixed: `NotImplementedError` with a message match
`prefill.py:1833` — reject fp8 unconditionally in single prefill	Declined; both arms already refuse it before bootstrapping

Copilot

🟡 Changes recommended

FP8 encoding validation can fail open, reused auto wrappers can retain stale backend selection, and two advertised page-size paths lack numerical coverage.

Get a fresh assessment by requesting another Copilot review.

Review details

Files reviewed: 16/16 changed files
Comments generated: 3
Review effort level: Balanced

**A wrapper that once resolved to fa2 refused fp8 for the rest of its life.** plan() went concrete on the first call and never re-read `_backend_requested`, so any first plan that hit an AITER constraint poisoned every later one -- and a served wrapper is re-planned every step. Reproduced: plan 1 (custom mask) backend = fa2 plan 2 (fp8) RAISED: NotImplementedError: fp8 prefill ... no in-tree fa2 kernel ... after: plan 2 (fp8) backend = aiter `auto` now re-resolves every plan, except under CUDA-graph capture, where the captured graph holds the first backend's buffers. **The encoding check failed open.** `_require_native_fp8_dtype` skipped itself when `_native_fp8_dtype()` returned None, and `_aiter_ops_importable()` only proves `aiter.ops` imports -- `aiter.dtypes` can be absent while the AITER backend is still selected. That is the NaN path, so it now refuses instead. **Page size is an AITER dispatch axis and only 16 was measured.** The routed set is {1, 16, 1024} and set membership is not coverage: an unsupported kernel at one page size is the exact failure this PR exists to fix. All three now run against the fp32 reference with a finiteness assertion, and all three pass. Co-Authored-By: Claude <noreply@anthropic.com>

Copilot

🔵 Needs a closer look

The FP8 path depends on pinned external AITER ABI details and architecture-specific GPU behavior requiring final human validation.

Review details

Files reviewed: 16/16 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>

…refill

Copilot

🟢 Approval recommended

The previously identified issues are addressed, and the implementation includes comprehensive routing, validation, and dual-architecture numerical coverage.

Review details

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

Copilot

🟡 Changes recommended

The descale paths lack independent coverage, and several user-facing FP8 behavior statements are inaccurate.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

docs/rocm/backends.md:708

This overstates the error contract for single prefill. With backend="aiter", single_prefill_with_kv_cache does not enter _reject_fp8_on_fa2; its existing AITER dtype validation raises RuntimeError, as the current implementation and prior verification show. Distinguish that case so the documentation does not promise NotImplementedError for every backend.

* Every other prefill route — single, ragged, and the paged flat-gather path —
  reaches `mha_fwd`/`mha_varlen_fwd`, which have no fp8 kernel. Those raise
  `NotImplementedError` naming fp8; the in-tree fa2 kernel rejects 8-bit types
  in a `static_assert`, which would otherwise surface as a compiler log.

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

Copilot round 5 (reviews 5192472001 and 5192482926). **The three descale pointers had no independent coverage.** Every test quantized K and V together and passed all three scales, so a swapped K/V pointer in the shim cancelled out and the "all three are required" guard was never exercised. `_fp8_kv_with_distinct_scales` now scales V 8x larger than K, which makes the swap observable, and the new cases are: swapping scale_k and scale_v must change the result; omitting any one of the three must raise; and passing descales with a bf16 query must raise. Seven descale tests, all green. Four user-facing claims were wrong, three of them mine: - The probe can demote a *routed* fp8 page size, and fp8 has no gather to fall back to -- it raises. The generic "falls back to the gather with a warning" was true only for fp16/bf16. - The encoding restriction `_require_native_fp8_dtype` enforces was documented nowhere the user would look, so picking the other advertised torch E4M3 dtype produced an unexplained rejection. - The constructor's `backend` docstring still said `auto` picks AITER for fp16/bf16 only, and `_auto_select_prefill_backend`'s said the same. - The fp8 section promised `NotImplementedError` on every non-paged route. Measured, single prefill under `backend="aiter"` raises `RuntimeError` from AITER's own dtype check -- which is exactly the evidence used to decline an earlier finding on that line, so the doc contradicted the decline. Co-Authored-By: Claude <noreply@anthropic.com>

…refill

demandal25 · 2026-09-13T22:19:18Z

Suppressed comment from review 5192482926 (docs/rocm/backends.md:708): fixed in 8d34104 — single prefill under backend="aiter" raises RuntimeError from AITER's own dtype check, not NotImplementedError, and the doc now says so. That is the same measurement used to decline the earlier finding on that route, so the doc had been contradicting the decline.

Copilot

🟡 Changes recommended

Distinct non-causal and soft-capped FP8 kernel variants remain untested, leaving dispatch-specific failures undetected.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

docs/rocm/backends.md:697

This introduces five bullets, not four. Update the count (or avoid a fixed count) so the fp8 limitations are documented consistently.

  matching upstream CUDA.

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

Copilot

🟡 Changes recommended

FP8 variants are missing from the prebuild model, and several routing diagnostics and benchmark assumptions remain inconsistent.

Get a fresh assessment by requesting another Copilot review.

Review details

Files reviewed: 16/16 changed files
Comments generated: 5
Review effort level: Balanced

…et count Copilot round 6. `needs_mask` and `has_logits_cap` each select a different fp8 `.so` filename in the loader, and every fp8 test used causal=True with no soft cap -- so a naming or bootstrap defect in the `_nmask` or `_logits` pertensor arms would have passed this suite and failed on a caller's first non-causal or soft-capped call. Probed all four combinations before writing anything, on gfx942: causal=True cap= 0.0 -> backend=aiter finite=True causal=False cap= 0.0 -> backend=aiter finite=True causal=True cap= 30.0 -> backend=aiter finite=True causal=False cap= 30.0 -> backend=aiter finite=True No defect, so this is coverage rather than a fix. `_plan_and_run` and `_reference` take `causal` and `logits_soft_cap`; the non-causal case is checked against the fp32 reference, and the capped cases against the *uncapped* result -- a dropped cap returns the uncapped answer, which a reference comparison with its own tanh model would be looser about. Also: the fp8 section said "Four constraints" and has five since the encoding restriction was added. The count is gone rather than corrected; it was only ever a thing to get wrong. Co-Authored-By: Claude <noreply@anthropic.com>

demandal25 · 2026-09-13T22:31:36Z

Suppressed comment from review 5192508979 (docs/rocm/backends.md:697): fixed in bb89eec — the count is removed rather than corrected, since it was only ever a thing to get wrong.

Copilot

🟢 Approval recommended

The implementation, guards, documentation, and targeted regression coverage are consistent, with prior findings addressed.

Review details

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

…contract Copilot round 7. All three are consequences of correcting the native set to {1, 16, 1024}, which this PR did and their text predates. **The standard paged-prefill benchmark rows measured the gather.** The testlist chose page 128 precisely to avoid the flat gather's in-timed-region copy, and 128 is not native -- it never was; that was the belief this PR disproved. Those six rows move to 1024, the one size `auto` routes natively for bf16. Verified they now reach the kernel the header claims: fa2 bs1 page1024 -> 0.2628 ms auto bs1 page1024 -> aiter 0.0547 ms fa2 bs16 page1024 -> 2.5992 ms auto bs16 page1024 -> aiter 0.5904 ms **The probe-failure warning promised a fallback fp8 never gets.** The text is shared, so an fp8 config whose native probe failed was told it was "falling back to the flat-gather path" immediately before plan() raised because no such kernel exists. The outcome clause is now dtype-aware, and the docstring's "falling back is always correct" says which dtypes it is correct for. **A route check used the wrong dtype.** The gather-fallback numerics test queried `_aiter_paged_route_page_sizes(torch.float16)` and then planned in bfloat16; the helper is dtype-dependent by design, so the two diverging would have made it run or skip on the wrong answer. Declined: adding the fp8/qscale axes to `flashinfer/jit/rocm/aiter_variants.py` so the prebuild store covers them. fp8 is batch-prefill only, and that family is already `servable_from_store=False` -- `_aiter_bootstrap_batch_prefill` is also the page-size capability probe, so it runs whether or not the .so is in the store. Prebuilding it fills the store with files nothing ever saves time on, which is what that flag exists to say. Co-Authored-By: Claude <noreply@anthropic.com>

Copilot

🟡 Changes recommended

FP8 scale arguments can be double-applied, and the routed-page probe-failure guard lacks direct regression coverage.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (2)

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

tests/rocm/test_fp8_paged_prefill.py:315

This bound is too loose to prove that K and V received the correct descales: scaling v by 8 also makes 40 * sv several units wide, so the internally swapped result can still satisfy it. The following not allclose only proves that the pointers affect different operands, not that they are wired in the right order. Compare both outputs against ref and require the declared ordering to have the lower error.

flashinfer/rocm/prefill.py:3195

The method now exposes two sets of FP8 calibration arguments. The existing q_scale/k_scale/v_scale parameters are still documented as FP8 scales and are applied to sm_scale/the output, while these new descales are also required and applied inside AITER. A caller following both parts of the API will therefore dequantize twice and get incorrect results; passing only the existing arguments fails the new required-descale check. Please make one set canonical for this route, or reject the legacy set for FP8 and explicitly document that it must not be combined with these arguments.

        scale_q : Optional[torch.Tensor]
            fp8 dequantisation scale for ``q``: a **per-tensor** float32 tensor of
            one element on ``q.device``. Required for an fp8 query, rejected
            otherwise. ``scale_k`` / ``scale_v`` are the same for the KV cache.
            A per-head tensor is silently read as element 0, so shape matters.

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

…ale order Copilot round 8. The first is a correctness bug this PR introduced. **Two calibration APIs on one call, and using both dequantizes twice.** `run()` already took float `q_scale`/`k_scale`/`v_scale`, which fold into `sm_scale` and the output; this PR added tensor `scale_q`/`scale_k`/`scale_v`, which AITER applies internally and which fp8 *requires*. A caller following both halves of the documented API got a silently wrong answer, and one following only the old half hit the required-descale check. The float set is now refused for an fp8 query, naming the tensor set, and the docstring says they are mutually exclusive and why. **The K/V ordering test did not test ordering.** It asserted `err < 40 * float(sv)` and `good != swapped`. Scaling V by 8 makes that bound several units wide, so the swapped result could satisfy it too, and "they differ" only shows the pointers reach different operands -- not that they reach the right ones. Both results are now compared against the fp32 reference and the declared order must have the lower error, which it does. **The probe-failure fp8 guard had no coverage.** Both non-routed-page tests exit at the earlier page-size check, so nothing exercised the guard inside the probe-failure branch -- the only thing standing between a routed page whose probe fails (an AITER source build is the usual cause) and the flat-gather path that has no fp8 kernel. Covered by forcing `_aiter_native_paging_available` false at a routed page size. Co-Authored-By: Claude <noreply@anthropic.com>

demandal25 · 2026-09-13T23:05:51Z

Suppressed comments from review 5192546694, both fixed in 67611e0.

prefill.py:3195 — the float q_scale/k_scale/v_scale are now refused for an fp8 query. They fold into sm_scale and the output while the tensor scale_* are applied by the kernel, so a caller following both halves of the API dequantized twice and got a silently wrong answer. The docstring says they are mutually exclusive.
test_fp8_paged_prefill.py:315 — correct, the bound was wide enough for the swapped result to pass and not allclose only proved the pointers reach different operands. Both results are now measured against the fp32 reference and the declared order must have the lower error.

Copilot

🔵 Needs a closer look

The hardware-specific AITER ABI, dynamic variant loading, and cross-architecture FP8 behavior warrant final human review.

Review details

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

…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>

demandal25 and others added 6 commits September 11, 2026 23:37

style(rocm): ruff-format the fp8 paged-prefill tests

fa95f76

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>

Copilot AI balanced review requested due to automatic review settings September 12, 2026 14:05

Copilot started reviewing on behalf of demandal25 September 12, 2026 14:06 View session

Copilot AI reviewed Sep 12, 2026

View reviewed changes

Comment thread benchmarks/rocm/bench_aiter_prefill.py

Comment thread flashinfer/rocm/prefill.py

Comment thread flashinfer/rocm/prefill.py Outdated

Comment thread flashinfer/rocm/prefill.py

Comment thread tests/rocm/test_fp8_paged_prefill.py Outdated

Comment thread docs/rocm/backends.md Outdated

This was referenced Sep 12, 2026

Route short-query paged prefill to fa2, at 1.19-6.08x over AITER #367

Closed

fp8 KV cache in the in-tree prefill kernel, and the refusals around it #368

Open

Merge branch 'amd-integration' into rocm-fp8-prefill

7bb6518

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

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