fp8 KV cache in the in-tree prefill kernel

[demandal25]

Summary

An fp16 query with an fp8 KV cache failed at JIT build on every in-tree prefill entry point — including BatchDecodeWithPagedKVCacheWrapper(use_tensor_cores=True), which the README advertises as supporting an fp8 KV cache. The failure was a ninja log, not an error. This makes the kernel serve fp8 KV, and makes every combination it still cannot serve fail with a sentence.

#366 (merged) covers fp8 on the AITER paged route, where q and kv must both be fp8. This covers the other shape — a 2-byte query against an fp8 cache, which is what an fp8 KV cache actually is — and it has no AITER route at all. Rebased onto amd-integration after #366 landed, so these commits stand alone.

What changed

Kernel

• include/flashinfer/rocm/attention/prefill.cuh — KernelTraits gains DTypeKVSmem. KV stays fp8 in HBM and is dequantized into a 2-byte LDS tile in the global→smem produce path, so DTypeKV is now the HBM type only. permuted_smem.cuh is untouched.

Refusals

• flashinfer/jit/rocm/modules.py — a dtype allowlist at the gen_customize_* seam for prefill and decode. q and o must be float16/bfloat16; KV may additionally be fnuz fp8. On prefill and POD a 2-byte KV must also equal the query; decode is exempt (see below). An explicit backend="aiter" with unequal q/kv dtypes is refused too. Every guard runs before its generator builds a URI or imports aiter, so an unmapped dtype gets the sentence rather than a KeyError, and the refusal still fires on an interpreter with no AITER wheel. The AITER allowlists are per-launcher: fp8 is admitted at the batch-prefill seam, which batch_prefill_paged_aiter.cu serves, and refused at the single-prefill and paged-decode seams, which are fp16/bf16 only.
• flashinfer/rocm/decode.py — both decode entry points validate k.dtype == v.dtype, sharing the prefill helper through the import edge decode already has.
• csrc/rocm/pytorch_extension_utils.h — CHECK_KV_DTYPES_MATCH, called from every attention run entry point (single_prefill.cu, both batch_prefill.cu runs, single_decode.cu, batch_decode.cu, pod.cu, batch_pod.cu). This is where the invariant lives: it covers the routes Python cannot reach without editing an upstream file, since both block-sparse wrappers call paged_run and POD is not shadowed under flashinfer/rocm/.
• include/flashinfer/rocm/attention/{pod,batch_pod}.cuh — a static_assert(!is_fp8_kv_v<...>). Both instantiate prefill's KernelTraits directly, and the deleted blanket assert was their only compile-time tripwire. The Python refusal sits in _gen_customize_pod_like_module, the one body both POD customize entry points delegate to, so it cannot be bypassed by calling the customize generator directly.
• flashinfer/rocm/prefill.py — _check_kv_dtypes_match on the two non-wrapper entry points, and the removal of a silent fp8-query downcast in the ragged wrapper.
• flashinfer/jit/utils.py — torch.float8_e5m2fnuz was missing from both dtype maps, so it raised KeyError from the URI builder.

Tests, docs, benchmark

• tests/rocm/test_fp8_kv_prefill.py — new, 99 cases.
• flashinfer/rocm/arch_caps.py, docs/rocm/backends.md, README.md (regenerated).
• benchmarks/rocm/bench_fa2_prefill.py — a --kv-dtype axis; _bytes() hardcoded 2 bytes for q, o, k and v alike, so every fp8 bandwidth number would have read ~2x high.

Architecture / design notes

Dequantize on load, not an fp8 LDS tile. Upstream keeps fp8 in LDS and dequantizes per fragment. Three measured facts say that costs a rewrite of the fragment-addressing layer and buys nothing here: NUM_MMA_KV is register-capped (ELEMS_PER_FRAGMENT = 16*16/64 = 4), so halving the KV smem type widens exactly one cell of the tile table; there is no async copy to forfeit, since commit_group/wait_group are no-ops and load_64b is a plain assignment, so every KV byte already passes through VGPRs; and the existing dead sizeof(DTypeKV)==1 branches are wrong, not merely unimplemented — they carry upstream's 8-half CUDA fragment geometry where a CDNA fragment is 4 halves. The seam back to an fp8 tile is the one conditional_t in kv_smem_type_t.

The "2-byte KV must equal the query" rule is prefill-only. Prefill instantiates the MFMA from DTypeQ and reads KV fragments as that type, so a bf16 cache under an fp16 query is misread. Decode cast_loads q and k into vec_t<float> separately and accumulates in float, so mixed 2-byte dtypes are correct there — measured: fp16-q with bf16-kv decode agrees with bf16/bf16 to 0.0024. Applying the rule to decode would have refused a working combination.

The k/v-match invariant is enforced in C++, with two Python calls that buy something extra. CHECK_KV_DTYPES_MATCH at every run entry point is the definition: block-sparse and POD reach the same kernels without passing through a shadowed module, and the alternative was editing flashinfer/sparse.py, which has no prior AMD commit and would carry a conflict hunk into every upstream sync. Python keeps the check only in single_prefill_with_kv_cache and single_decode_with_kv_cache, where it runs before module resolution and so refuses without paying a cold JIT. The three wrapper run() copies were dropped: plan() has already built the module there, so they added 335 ns and a second spelling of the same rule. Mismatches from a wrapper therefore surface as RuntimeError, not ValueError.

kv_smem_type_t keys on the fp8 types, not on sizeof()==1. An int8/uint8 cache has no float interpretation; widening it would route through vec_cast's primary template, which reads 0..255 as a value. Keying on size would have converted the old blanket static_assert from a refusal into silent garbage.

Dequant goes through vec_cast<float, fp8>. vec_cast<half, fp8> at vec_size >= 2 decodes fnuz with the OCP bias (- 7 where the format's bias is 8) and mishandles denormals — a latent 2x error in shared code. It stays unreachable, and the float route is what working decode already exercises. Left unfixed deliberately: nothing under tests/rocm/ can call vec_cast today, so a conformance test needs a new .cu, a JIT generator and a binding.

The OCP spellings are refused rather than mapped. dtype_map_hip sends torch.float8_e4m3fn and torch.float8_e5m2 onto the fnuz C++ types, whose exponent bias is one greater (E4M3 7 against 8, E5M2 15 against 16). The refusal is keyed on backend != "aiter", not == "fa2", because ROCm logs and ignores an explicit "fa3" and routes it to the same in-tree kernel. AITER keeps its own fp8 dtypes — it compiles against aiter.dtypes.fp8, which is OCP e4m3fn on gfx950.

fp8 KV stays JIT-only. rocm/aot.py draws dtype_kv from the f16 list; adding a KV-dtype axis multiplies the prefill+decode matrix and is a wheel-size decision deserving its own measurement. First use pays a cold build.

Benchmark results

Prefill: fp8 is neutral, and that is the honest result. These shapes are compute-bound — arithmetic intensity 205–3277 FLOPs/B against MI300X's ~247 ridge point — so halving KV bytes buys nothing. gfx942, llama3-8b (GQA 32/8, hd 128), causal, median ms:

seq	fp16	fp8 e4m3fnuz
512	0.074	0.074
1024	0.239	0.234
2048	0.744	0.751
4096	2.277	2.319
8192	8.708	8.322

No fp16 regression, which is the point of produce_kv_unit forwarding verbatim for a 2-byte cache. Same shapes, kernel reverted vs applied: 0.074 / 0.240 / 0.744 / 2.263 / 8.701 → 0.074 / 0.239 / 0.744 / 2.277 / 8.708 ms, inside run-to-run noise.

Decode is where an fp8 cache pays, and the comparison that matters is against the only fp8 decode route that previously worked. gfx942, GQA 32/8, hd 128, page 16, median ms:

bs / kv_len	plain decode, fp8	tensor-core decode, fp8	speedup
1 / 8192	0.083	0.045	1.84x
8 / 4096	0.212	0.110	1.93x
32 / 2048	0.387	0.198	1.95x
64 / 1024	0.386	0.192	2.01x

Against tensor-core fp16 the same rows are only 1.02–1.17x (mean 1.09x): at 1.2 TB/s of MI300X's 5.3 the path is latency-bound, not bandwidth-bound. And plain decode with fp16 remains the fastest route overall (0.041–0.106 ms) — an fp8 KV cache is a memory-capacity trade, not a speed win against fp16. Two things worth flagging for follow-up, neither touched here: plain decode with fp8 is 2–3.6x slower than plain decode with fp16, and the tensor-core decode path leaves most of the memory system idle.

Test plan

• Equivalence oracle — fp8 against the torch-dequantized fp16 run, asserted bitwise (torch.equal, no tolerance): single / paged / ragged / tensor-core decode x 2 fp8 dtypes x head_dim 64,128, on gfx942 and gfx950.
• Block-sparse fp8 KV — bitwise identical to the dequantized reference for both fnuz dtypes; both wrappers build through the batch-prefill module.
• k_scale/v_scale on an fp8 cache — asserted as the two folds the implementation performs (k_scale into sm_scale, v_scale linear on the output). A/B: disabling sm_scale *= k_scale fails all four cases.
• Refusals — fp8 query, OCP KV, integer KV, float32, mismatched k/v on prefill, decode and block-sparse, POD + fp8, and the AITER generator allowlist (10 dtype combinations x 2 generators, each seam asserted separately since fp8 is legal at one and not the other).
• Every public generator refuses before building its URI — swept over all 13 gen_*_module entry points, with the list derived from dir(modules) so the test fails if a new generator is added without one. The hand-written list previously covered 6 of 13, which is how the two customize-POD generators kept building fp8 specs the POD kernel then rejected in ninja.
• The C++ backstop is live, not merely written: with the Python check removed, block-sparse with k=fp16, v=bf16 raises RuntimeError: paged_k_cache has dtype Half but paged_v_cache has BFloat16.
• Reproduced the reported defect at 6dc0b0fc0 first — 54 / 54 / 27 compiler error lines from the three entry points — and the k/v hazard: k=e4m3fnuz, v=fp16 returned finite=False with no error.
• Regression, gfx942: the prefill/POD family (test_{single,batch}_prefill_kernels, bf16_custom_mask, sliding_window, pod, batch_pod, customize_prefill_*, prefill_decode_dispatch, fp8_paged_prefill) — 8,627 passed, 3,563 skipped, 0 failed.
• Regression, gfx942, after the AITER allowlist: test_fp8_paged_prefill (fp8 paged prefill on the AITER route, at 1.22-1.66x over bf16 #366's) + test_batch_prefill_kernels — 5,615 passed, 2,136 skipped, 0 failed.
• Final tree, gfx942 (9be6d0960): test_fp8_kv_prefill — 119 passed, 16 skipped, 0 failed; test_block_sparse, test_pod, test_batch_pod, test_batch_decode_kernels green on the preceding commit.
• gfx950: test_fp8_kv_prefill, test_fp8_paged_prefill, test_batch_decode_kernels, test_single_prefill_kernels. Two failures, both pre-existing — test_gated_architecture_really_is_defective[17-2048] and [512-512] fail identically on the base 6dc0b0fc0, same GPU and image.
• gfx950 caught a CDNA4-only build break this series had introduced and gfx942 could not see: staging the fp8 load in a DTypeKV[GRANULE] does not compile where HIP_FP8_TYPE_FNUZ=0 leaves the fnuz types without a default constructor. Fixed in 54cdcaecf by staging in a uint32_t, which keeps the load aliasing-clean; re-verified on MI350X, then re-run green on gfx942. Neither conformance nor pre-commit compiles kernels, so only an on-device build surfaces this class of defect.
• Rebased onto amd-integration after fp8 paged prefill on the AITER route, at 1.22-1.66x over bf16 #366, Correct the asm prefill speedup, which the old benchmark overstated #370, Let the soft-cap gate test reach the kernel it measures #371 and Record why ragged prefill is ungated, and what would have to be measured to gate it #372 merged. The benchmark numbers above predate that rebase; none of those merges touch the fa2 prefill kernel.
• pre-commit run -a.

Known limitations

• In-kernel RoPE prefill is non-deterministic, for every KV dtype. pos_encoding_mode="ROPE_LLAMA" on the fa2 prefill kernel returns a different answer on each call from identical inputs — measured on an unmodified kernel with plain fp16 KV (spread 0.2-1.4 over repeats, 6 runs giving 6 distinct results, already at a single KV tile). Pre-existing and dtype-independent, so this series does not gate fp8 on it: that would imply the fp16 path is sound. No prefill test covers RoPE, which is how it survived. Recorded in docs/rocm/backends.md; fixing it needs its own change.
• fp8 halves bytes moved, not load instructions. The produce path issues one 4-byte global load per lane where a 2-byte cache issues 8, and the trip count is unchanged, so prefill measures flat. Claiming the bandwidth win needs an 8-element-per-lane load writing two swizzled slots — a separate change.
• fp8 KV is JIT-only. rocm/aot.py draws dtype_kv from the f16 list; adding a KV-dtype axis multiplies the prefill+decode matrix and is a wheel-size decision deserving its own measurement. First use pays a cold build.

[Copilot]

🟡 Changes recommended

Remaining dtype-validation gaps can permit unsupported mixed operands or silently reinterpret decode cache data.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds FNUZ FP8 KV-cache support to ROCm’s in-tree prefill kernel while providing explicit failures for unsupported dtype combinations.

Changes:

• Dequantizes FP8 KV data into 2-byte LDS tiles for prefill and tensor-core decode.
• Adds dtype validation, AITER FP8 paged-prefill routing, and regression tests.
• Updates capability documentation and benchmark dtype/byte accounting.

File summaries

File	Description
include/flashinfer/rocm/attention/prefill.cuh	Implements FP8 KV loading and LDS conversion.
flashinfer/jit/rocm/modules.py	Adds ROCm attention dtype guards.
flashinfer/rocm/prefill.py	Adds routing, validation, descales, and output handling.
flashinfer/jit/utils.py	Maps E5M2FNUZ for JIT generation.
tests/rocm/test_fp8_kv_prefill.py	Tests in-tree FP8 KV behavior and refusals.
tests/rocm/test_fp8_paged_prefill.py	Tests AITER FP8 paged prefill.
tests/rocm/test_prefill_decode_dispatch.py	Updates native page-size expectations.
tests/rocm/test_batch_prefill_kernels.py	Updates AITER routing tests.
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 for FP8.
csrc/rocm/batch_prefill_paged_aiter.cu	Validates and launches FP8 paged prefill.
csrc/rocm/batch_prefill_paged_aiter_jit_pybind.cu	Extends the paged binding signature.
csrc/rocm/batch_prefill_aiter_customize_config.jinja	Adds required HIP dtype headers.
csrc/rocm/aiter_loader.cc	Resolves scaled FP8 AITER variants.
flashinfer/rocm/arch_caps.py	Updates advertised capabilities.
docs/rocm/backends.md	Documents FP8 support and limitations.
README.md	Regenerates the backend capability table.
benchmarks/rocm/bench_fa2_prefill.py	Adds FP8 KV benchmarking and byte accounting.
benchmarks/rocm/bench_aiter_prefill.py	Aligns benchmark page sizes with routing.

Review details

Suppressed comments (1)

flashinfer/jit/rocm/modules.py:463

• The exported flashinfer.jit.gen_customize_batch_decode_module bypasses this guard: its implementation at modules.py:841-919 feeds the supplied dtype directly into dtype_map_hip. Thus OCP fp8 remains reachable through a supported public JIT API and is silently interpreted as FNUZ. Apply the check in the custom batch generator, letting this wrapper inherit it.

    _check_fa2_fp8_dtypes("batch decode", dtype_q, dtype_kv, dtype_o)

• Files reviewed: 19/19 changed files
• Comments generated: 3
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[demandal25]

Suppressed comment (review of 2026-09-12T23:22Z), modules.py:463 — gen_customize_batch_decode_module bypassing the guard: fixed in 42bdd81, along with the single-decode twin.

[Copilot]

🟡 Changes recommended

Auto backend selection remains sticky across replans, and the new scale API conflicts with existing signatures and documented scaling parameters.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 23/23 changed files
• Comments generated: 4
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[demandal25]

Disposition log for the bot reviews on this PR, at f50c83467. Inline findings are answered in their own threads; this covers the suppressed ones, which have no thread, plus the standing declines.

Correction to an earlier verdict. I previously declined the suppressed comment on modules.py:990 (AITER dtype checking) on the float32 half alone, and never tested the int8 half it also named. That half was real and worse than a slow error — there was no error at all:

aiter  torch.int8   -> NO RAISE, spec=batch_prefill_with_kv_cache_aiter_dtype_q_i8_...
aiter  torch.uint8  -> NO RAISE, spec=batch_prefill_with_kv_cache_aiter_dtype_q_u8_...

Both are in dtype_map_hip, so they passed the equality check, built a URI and produced a JitSpec headed for ninja with no kernel behind it. Accepted and fixed in 719560c2a: an AITER allowlist (float16, bfloat16, e4m3 in both arch spellings) placed before URI construction in all four generators, verified over 10 dtype combinations x 2 generators, with #366's suite re-run green (5,615 passed, 2,136 skipped, 0 failed).

Suppressed comments, per review:

• modules.py:463 (review 1) — answered in this comment.
• sparse.py:1631 (review 3), planned-vs-runtime dtypes — moot: flashinfer/sparse.py is no longer touched by this PR (git diff origin/amd-integration..HEAD -- flashinfer/sparse.py is empty). The k/v check moved to CHECK_KV_DTYPES_MATCH at the C++ entry points, which covers both block-sparse wrappers via paged_run. The gap it names is pre-existing: git show origin/amd-integration:flashinfer/sparse.py | grep -c _check_cached_qkv_data_type -> 0.
• modules.py:990 (review 6) — accepted, see the correction above.
• modules.py:774 (review 7) — accepted, same finding as that review's inline comment on :136: the AITER output allowlist admitted any 2-byte dtype, while the launchers pin it exactly (single_prefill_aiter.cu:85 and batch_ragged_prefill_aiter.cu:56 require o == q_dtype; batch_prefill_paged_aiter.cu:67 requires bf16 for an fp8 query). fp16 q/kv with bf16 output passed, paid a cold JIT, then failed in TORCH_CHECK. Fixed in 7a8824749, verified over 6 output combinations.

Standing declines, so later rounds re-raising them can be read against this:

• 4-byte staged load — the 2-byte path already casts the same pointer to b64_t*, which requires 8-byte alignment, so the codebase assumes strictly more than this does.
• torch.float8_e5m2fnuz in filename_safe_dtype_map — that dict already carries torch.float8_e4m3fnuz, an AMD-only dtype, so the ROCm precedent is established in that exact dict.
• RoPE + fp8 — the non-determinism is dtype-independent and hits fp16 today, so refusing only fp8 would imply the fp16 path is sound. Recorded under Known limitations.
• POD guard on gen_pod_module/gen_batch_pod_module rather than the gen_customize_* seam — gen_customize_pod_module is not re-exported from flashinfer/jit/__init__.py (only gen_customize_batch_decode_module is, which is why the decode guard moved down), and its only callers are those two guarded functions.
• POD cross-phase dtype validation — flashinfer/pod.py already calls _check_cached_qkv_data_type for both phases against the same cached plan dtypes, so the split described cannot be reached through the public wrapper.
• AITER bootstrap ordering — the cost is a slow refusal on an already-invalid call, not a wrong result, and re-ordering fp8 paged prefill on the AITER route, at 1.22-1.66x over bf16 #366's bootstrap is what broke its fp8 prefill once already here.
• single_decode.cu missing the k/v check — it was never missing; line 50 enforced it as CHECK_EQ. Now spelled CHECK_KV_DTYPES_MATCH so all seams share one message.

Follow-ups for amd-integration, not this PR: the sticky-auto re-plan on both wrappers, and the scale_q/scale_k/scale_v overload and doc surface (both from review 2, both blame to #366). Plus one performance item from the self-review: the fp8 produce path issues a 4-byte load per lane where a 2-byte cache issues 8, so instruction count does not fall with byte count — plausibly why prefill measures flat.

Note this branch is maintained with fixups and autosquash, so Fixed in <sha> in older threads may point at a commit a later rebase superseded. The series is 518806385, 719560c2a, 8e761617d, 2835a84a9, 4fd9e14c8, then the amd-integration merge, then 7a8824749 and f50c83467.

[Copilot]

🟡 Changes recommended

Sparse wrappers can still silently reinterpret runtime KV tensors, and some unsupported dtypes bypass the intended validation.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

flashinfer/sparse.py:1631

• The K/V equality check still permits both tensors to differ from the dtype used to build the cached module. Planning this wrapper for fp16 and running it with two fp8 tensors therefore passes here and lets paged_run reinterpret both buffers as fp16. Persist the planned Q/KV dtypes and validate runtime inputs against them before reshaping and dispatch.

        if k.dtype != v.dtype:
            raise ValueError(
                f"k has dtype {k.dtype} but v has {v.dtype}; block-sparse "
                "attention compiles a single KV dtype and would reinterpret v."
            )

• Files reviewed: 13/13 changed files
• Comments generated: 2
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[Copilot]

🟡 Changes recommended

POD validation still permits cross-phase dtype mismatches that can reinterpret query and KV tensor bytes.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 18/18 changed files
• Comments generated: 6
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[Copilot]

🟡 Changes recommended

The FP8 conversion has an alignment hazard, and one unsupported dtype combination still raises a bare assertion.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 18/18 changed files
• Comments generated: 2
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[Copilot]

🟡 Changes recommended

AITER dtype validation remains incomplete and can occur only after expensive bootstrap work or URI construction.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

flashinfer/jit/rocm/modules.py:990

• Equality is not a sufficient AITER dtype check here. For example, dtype_q=dtype_kv=dtype_o=torch.float32 passes this branch, while gen_batch_prefill_module constructs the URI first and raises a bare KeyError because float32 is absent from filename_safe_dtype_map; mapped but unsupported types such as int8 can proceed even further. Add an AITER-specific allowlist before URI construction so unsupported equal-dtype requests also receive the promised typed refusal.

    if backend == "aiter" and dtype_q != dtype_kv:
        raise NotImplementedError(
            f"batch prefill: AITER requires equal query and KV dtypes; got "
            f"{dtype_q} and {dtype_kv}."
        )

• Files reviewed: 18/18 changed files
• Comments generated: 1
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[Copilot]

🟡 Changes recommended

The AITER allowlist admits output and single-prefill FP8 combinations that the generated launchers cannot execute.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

flashinfer/jit/rocm/modules.py:774

• This uses the batch-paged dtype policy for single prefill and therefore admits e4m3 fp8 Q/KV. However, the generated single_prefill_aiter.cu launcher explicitly accepts only fp16/bf16 at lines 81-84, so both this public custom generator and the regular generator that delegates here can create a module that can never execute. Reject _AITER_FP8_DTYPES at the single-prefill seam; fp8 AITER support is paged batch only, and update the test that currently marks these single-generator cases as allowed.

    if backend == "aiter":
        _check_aiter_dtypes("single prefill", dtype_q, dtype_kv, dtype_o)

• Files reviewed: 19/19 changed files
• Comments generated: 1
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[Copilot]

🟡 Changes recommended

The new single-decode dtype guard lacks a direct regression test.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 19/19 changed files
• Comments generated: 1
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[Copilot]

🟡 Changes recommended

The exported decode JIT builders can still raise opaque KeyError exceptions before the new dtype validation runs.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

flashinfer/jit/rocm/modules.py:938

• The exported flashinfer.jit.gen_batch_decode_module still calls get_batch_decode_uri before this validation (modules.py:519-530). Consequently an unmapped unsupported dtype such as torch.float32 raises KeyError from that public builder rather than the promised actionable refusal. Add the same check at the beginning of gen_batch_decode_module, before URI construction, and keep this customize-level backstop.

    _check_fa2_fp8_dtypes(
        "batch decode", dtype_q, dtype_kv, dtype_o, kv_must_match_q=False
    )

• Files reviewed: 19/19 changed files
• Comments generated: 1
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[demandal25]

Suppressed comments, reviews through f50c83467:

• modules.py:938 (batch decode builds its URI before validating) — Fixed in 34262802f, together with the inline twin at :681.
• modules.py:774 (single-prefill AITER admits fp8 its launcher rejects) — Fixed in c63616a72. Correct, and further than stated: batch_ragged_prefill_aiter.cu is fp16/bf16 only too, so fp8 is now allowed at the batch seam alone, which is the only one batch_prefill_paged_aiter.cu serves.

Also fixed from an internal review of the same push: gen_customize_pod_module and gen_customize_batch_pod_module bypassed the POD fp8 refusal entirely (9be6d0960), and gfx950 could not compile the series at all (54cdcaecf) — HIP_FP8_TYPE_FNUZ=0 leaves the fnuz types without a default constructor, so the DTypeKV[GRANULE] staging array was a CDNA4-only build break that gfx942 cannot see.

[Copilot]

🔵 Needs a closer look

Low-level LDS layout, pointer arithmetic, and cross-architecture FP8 behavior require final human GPU validation.

Review details

• Files reviewed: 19/19 changed files
• Comments generated: 0 new
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[demandal25]

Closing the review loop: the review on 9be6d0960 came back with no inline comments and no suppressed block, and CI is green (conformance, pre-commit).

Deliberately not in this PR, each recorded under Known limitations in the description rather than carried silently: the RoPE prefill non-determinism (pre-existing and dtype-independent), the flat prefill throughput (fp8 halves bytes moved but not load instructions — a wider per-lane load is its own change), and AOT packaging of an fp8 KV axis. Later rounds re-raising these should be read against this comment.