[CK] Add new FMHA batch prefill kernel with FP8 per-tensor and per-block KV quantization on gfx950

[poyenc]

Summary

V3 FMHA pipeline for batch prefill with paged KV cache on gfx950 (FP8 only). Uses the same 8-warp, 256x64 tile, double-buffered LDS architecture as the contiguous fmha_fwd V3 kernel, extended with scatter-gather page table support.

Supported configurations:

• FP8 per-tensor and KV block-scale quantization
• SGLang (1D) and vLLM (2D) page tables
• LINEAR KV layout only (VECTORIZED falls back to V2 via trait matching)
• No mask / causal mask
• Page sizes: 1, 16, 1024

Performance vs V2 (FP8 batch prefill, paged KV, stock ROCm 7.1.1, avg runs 2–6, MI355X):

Per-tensor (ps=1, sglang):

Problem	V2 TFlops	V3 TFlops	V3/V2
h=6/1 sq=1k c	43.1	47.7	1.11x
h=6/1 sq=2k c	95.4	120.1	1.26x
h=6/1 sq=4k c	197.5	259.5	1.31x
h=6/1 sq=8k c	396.3	549.4	1.39x
h=6/1 sq=16k c	778.7	1066.1	1.37x
h=6/1 sq=32k c	1000.2	1282.3	1.28x
h=6/1 sq=65k c	1045.0	1375.5	1.32x
h=6/1 sq=131k c	1072.9	1390.2	1.30x
h=16/1 sq=65k c	1072.3	1381.8	1.29x
h=40/40 sq=37k nc	930.1	1304.3	1.40x

KV block-scale (ps=1024, sglang):

Problem	V2 TFlops	V3 TFlops	V3/V2
h=6/1 sq=1k c	42.6	45.2	1.06x
h=6/1 sq=2k c	95.2	111.6	1.17x
h=6/1 sq=4k c	197.8	241.1	1.22x
h=6/1 sq=8k c	389.9	509.7	1.31x
h=6/1 sq=16k c	745.6	1013.1	1.36x
h=6/1 sq=32k c	1016.1	1205.4	1.19x
h=6/1 sq=65k c	1057.0	1313.5	1.24x
h=6/1 sq=131k c	1069.6	1330.3	1.24x
h=16/1 sq=65k c	1084.4	1324.2	1.22x
h=40/40 sq=37k nc	945.0	1251.6	1.32x

Key design decisions

Branchless page_id clamping — The scatter-gather path bypasses buffer descriptor NUM_RECORDS protection, so auto-advancing past seqlen_k causes XNACK faults. Solved with min(page_id, max_page_table_idx) instead of branch guards, which would fragment sched_group_barrier-scheduled basic blocks.

Page advance issue/consume split — Page table lookups (global_load_dword) are issued BEFORE buffer_loads so they sit oldest in the vmcnt FIFO. At consume, s_waitcnt(N) drains only the page lookup while keeping N buffer_loads in flight. sched_barrier(0) prevents reordering.

KV block-scale k_descale fold — k_descale is folded into the scalar row_max via the FP8 shift trick, eliminating a full-tile VALU pass. v_descale is merged into the softmax rescale factor (o_acc_scale), with the final correction applied in the epilogue.

No no-packed-fp32-ops kernel attribute — V3 does NOT use kernel_attr<true> (which adds target("no-packed-fp32-ops")) despite it giving a +2–4% benefit on per-tensor workloads. The attribute conflicts with explicit v_pk_mul_f32 inline asm used in the KV block-scale descale path (pk_mul_f32 helper): when the asm is inlined into the attributed kernel entry, the assembler rejects the instruction. Benchmarks show the attribute is neutral-to-positive for block-scale (removing it actually gives +1–3%), so we accept the 2–4% per-tensor trade-off to avoid the asm conflict. A future per-variant solution (attribute on for per-tensor, off for block-scale) could recover this, but the codegen currently uses a single F_kernel_attr per kernel.

New files

File	Description
fmha_batch_prefill_v3_kernel.hpp	V3 kernel with paged KV DRAM views, page table Kargs, batch/group mode
block_fmha_batch_prefill_v3_pipeline.hpp	V3 pipeline with scatter-gather loads, KV block-scale support
block_fmha_batch_prefill_v3_pipeline_default_policy.hpp	FP8 LDS descriptors, SRD rebasing for large page sizes
block_fmha_fwd_v3_detail.hpp	Shared V3 helpers: CoreLoopSchedulingParams, VALU intrinsics, macros

Codegen changes (fmha_batch_prefill.py)

Refactored to support per-architecture kernel generation (matching fmha_fwd.py patterns):

• ArchTrait / factory hierarchy: gfx9 -> gfx950 with arch-specific tiles, pipelines, and compatibility rules
• Separate fmha_batch_prefill_v2() / fmha_batch_prefill_v3() dispatch functions (V3 tried first, falls back to V2)
• 80 V3 kernels generated for gfx950

Bug fix in fmha_fwd.py

Fixed missing return in check_hdim — bare False was a no-op, causing unnecessary bias/dropout kernels for hdim=(192,128).

Dependencies

Based on #4437 (users/poyenc/fa-v3-fp8-pertensor). That PR must be merged first.

Test plan

• Full test_batch_prefill.py suite

[poyenc]

still be blocked by unrelated compilation errors

[poyenc]

the CI always failed to pull CK image

[Jeff-Huang]

To fix the OOB issue surfaced by AICK-1171, I based the V2 changes on your PR and opened a separate PR — #6932 — for the V2 load_physical_pages part only. Two small tweaks while extracting:

• max_page_table_idx is mandatory (no INT32_MAX default), so every callsite has to pass the bound explicitly. With the optional default, any unupdated callsite silently no-ops the clamp.
• Clamp applied to all branches in load_physical_pages (K prefetch + V LINEAR + V crosses-pages + V single-page lane0 broadcast) in V2.

Verified on MI-308X with the AICK-1171 reproducer and the full FMHA batch prefill suite on gfx942/gfx950. Splitting it out lets the OOB fix land on its own timeline; the rest of #6054 stays as-is.

Credit to your original approach — same min(page_id, max_bound) idea, just hardened.

[Jeff-Huang]

Hi @poyenc — thanks for the V3 batch_prefill kernel. While reviewing this PR I hit a silent-correctness issue in the new V3 KV_BLOCKSCALE arm. The output shape and magnitude look right (out_norm ≈ ref_norm within ~0.5%), so i
t doesn't crash — it just quietly exceeds the FP8 reference threshold at specific multi-page configurations. I tracked the root cause down and verified a minimal fix; sharing the full picture below.

Reproducer

A targeted regression test is in aiter PR #3032 (branch jeff-huang/fix-batch-prefill-oob-page-table-read).

test_batch_prefill_4gb_boundary_targeted is a small 2-page boundary probe designed to surface > 4 GB KV-cache page-table bugs without dilution. It allocates a 5000-page paged-KV pool (~5 GB for fp8, ~10 GB for bf16) with page_ size=1024 / num_kv_heads=8 / hdim=128, then for the kernel call passes a single-batch, 2-page page table [vh_start, vh_start + 1] with kv_len = 2048. Because only 2 pages are read and kv_len is small, ALL attention math ru
ns through the suspect pages — wrong reads cannot be diluted by surrounding correct pages (which is what hides the bug in the existing test_batch_prefill_large_kvcache long-sequence stress test).

The failing pytest parametrize id is [False-2-kv_blockscale-fp8]:

Param	Value	Meaning
causal	False	Plain attention, no causal mask
page_offset_factor	2	Place first page at the 2 × 2³¹ = 2³² byte boundary
quant_mode	kv_blockscale	Per-page FP32 descale (one descale per page, per K/V)
input_dtype	fp8	KV stored as 1 byte/element

For fp8, each page is 1024 × 8 × 128 × 1 byte = 1 MB. With page_offset_factor=2, the test picks page_indices = [4096, 4097] so the first page byte_offset is exactly 2³² (4 GB).

Note: the 2³¹ / 2³² byte boundary is incidental to this bug — test_batch_prefill_4gb_boundary_targeted was originally designed to surface a separate > 4 GB overflow concern, but happens to also trigger the bug des
cribed here. The actual root cause has nothing to do with byte offsets crossing 32-bit boundaries — any cross-page transition with KV_BLOCKSCALE on V3 fp8 will trigger the same hazard. This test case just happens to be a dete
rministic, low-noise way to expose it.

After checking out aiter PR #3032 with this PR's CK:

HIP_VISIBLE_DEVICES=0 python3 -m pytest \
  op_tests/test_batch_prefill.py::test_batch_prefill_4gb_boundary_targeted \
  -k "False-2-kv_blockscale-fp8" -v

Expected:

FAILED test_batch_prefill_4gb_boundary_targeted[False-2-kv_blockscale-fp8]
AssertionError: FP8 kernel vs reference difference too large:
                0.080078125 (threshold: 0.055)

Confirmed via stream_config.log_level_=1 that the failing case dispatches to ..._v3_..._kv_blockscale_..._linear_sglang.

Bug fingerprint (all conditions must hold)

Isolated by 23 page-permutation probes:

Condition	Required
arch = gfx950 + ROCm 7.1+	Yes — gfx942 with same source PASSes (V3 not built → V2 fallback)
quant_mode = KV_BLOCKSCALE	Yes — V3 PERTENSOR with same shape PASSes (max = 0.005)
kHasLogitsSoftCap = false	Yes — activates the kFoldKDescale path
causal = false	Yes
Multi-page (≥ 2 distinct pages)	Yes
Both pages ≥ overflow_page (= 2048 in this config)	Yes
Pages consecutive (gap = 1)	Yes
First page index is even	Yes (empirical filter — see note below)

Representative probe results — each row's bracket list is the physical page indices passed via kv_page_indices to the kernel. [4096] means a single page (page index 4096 in the paged-KV pool); [4096, 4097] means two pages (4096 then 4097). Note that each individual page reads correctly, only the transition fails:

pages          verdict   max_diff   note
[4096]         PASS      0.0054     single page (no transition)         ← baseline noise
[4097]         PASS      0.0063     single page (no transition)         ← baseline noise
[4096, 4096]   PASS      0.0054     duplicate (no transition)           ← baseline noise
[4096, 4097]   FAIL      0.0801     even-first consecutive — bug fires, OVER threshold
[4098, 4099]   FAIL      0.1108     even-first consecutive — bug fires, OVER threshold
[4097, 4098]   PASS      0.0313     odd-first — bug fires, ~6× baseline, just under threshold
[4096, 4098]   PASS      0.0157     gap=2     — bug fires, ~3× baseline, just under threshold
[4095, 4096]   PASS      0.0247     boundary  — bug fires, ~5× baseline, just under threshold

The "PASS" rows in the multi-page section also have elevated max_diff (3-6× the single-page baseline of ~0.005). All of them are silently affected by the same hazard — they only PASS because the random per-page descale values happen to differ less, keeping the multiplicative-factor error under the 0.055 FP8 threshold. With a different RNG seed or different head_dim/page_size, the same configurations could easily flip to FAIL. The "first page index even" condition in the fingerprint table is therefore an empirical filter for "loud enough to fail under this seed", not a fundamental requirement of the bug. Verification (further down) shows the fix drops every multi-page row back to baseline noise.

Root cause — read-after-write hazard between save_descales and fmha_alu0

The kFoldKDescale path (active for KV_BLOCKSCALE && !kHasLogitsSoftCap) has a 2-slot double-buffer read-after-write hazard within a single core_loop iteration.

All line numbers below are verified against PR HEAD 70ed05b674.

Notation: K_n = the specific K-tile that gemm0 of THIS iteration consumes; K_{n+1}, K_{n+2} = subsequent K-tiles in the kv_len sequence (consumed by later iterations). pi is the compile-time iteration index that alternates between 0 and 1 per core_loop call — see L1206 (iteration = [&](auto pi)) and L1411 (return iteration(number<0>{}) && iteration(number<1>{})).

Slot binding: at L1210 auto K_w0_lds_wr_idx = number<1>{} - pi; so K_w0_lds_wr_idx always evaluates to slot 1-pi. Both save_descales (writer) and fmha_alu0 (reader inside cl_calc(gemm1)) target this same slot in the SAME iteration.

Within one iteration(pi) (Wave0-3 region at L1259-1295; Wave4-7 region at L1357-1408 is structurally identical with K_w4_lds_wr_idx):

iteration(pi):                                        ┌── slot 1-pi state ──┐
  ① cl_calc(p01, gemm0)                               │  desc_for(K_n)       │  consumes K_n (loaded 2 iters ago)
  ② fmha_alu1(p23)                                    │  desc_for(K_n)       │
  ③ fmha_logits_trans(p01)                            │  desc_for(K_n)       │
  ④ save_descales(K_w0_lds_wr_idx)  // = slot 1-pi    │  desc_for(K_{n+2})   │  ⚠️ OVERWRITES K_n's descale
  ⑤ K_page_issue                                      │                      │
  ⑥ cl_load(memK, K_w0_lds_wr_idx, V_w0_lds_rd_idx)   │  K_{n+2} → LDS slot 1-pi │
  ⑦ K_page_consume                                    │                      │
  ⑧ cl_calc(p23, gemm1)  // calls fmha_alu0(1-pi)     │  reads desc_for(K_{n+2})│  ❌ but applies to K_n!
  ⑨ fmha_alu_D_upd_unpack(p23)                        │                      │
  ⑩ cl_load(memV, V_w0_lds_wr_idx, K_w0_lds_rd_idx)   │  K_lds_load(pi) → kv_tile.k_tile for next iter │

save_descales at step ④ writes the descale for the K-tile being LOADED THIS iter (K_{n+2}) into slot 1-pi. fmha_alu0 at step ⑧ reads the SAME slot 1-pi, expecting the descale of the K-tile that gemm0 just CONSUMED THIS iter (K_n).

Pipeline lag derivation — K_n ≠ K_{n+2} because cl_load(memK) writes the LDS slot in step ⑥ but kv_tile.k_tile in step ① was populated by K_lds_load in step ⑩ of an EARLIER iteration. Tracing through pre-stage + 2 iterations of core_loop: the K-tile cl_load'd into LDS at iteration i is consumed by gemm0 at iteration i+2 (the toggling of K_w0_lds_wr_idx = 1 - pi and K_w0_lds_rd_idx = pi swaps roles per iteration). The 2-slot buffer is too small to hold both the about-to-be-loaded descale AND the about-to-be-consumed descale simultaneously, and the current ordering puts the WRITE (step ④) before the READ (step ⑧) in the same iteration — closing the hazard.

Concrete trace at the page boundary for [4096, 4097] (d_a = page-4096 desc = 0.01056, d_b = page-4097 desc = 0.01038):

step iter consumed K  saved_k_descale[1] before/after save  fmha_alu0 reads → applies to
─────────────────────────────────────────────────────────────────────────────────────────
step7  0  K13         d_a → d_a   d_a → K13 ✓
step7  1  K14         (writes slot 0 — slot 1 untouched)
step8  0  K15         d_a → d_b   d_b → K15 ❌  (page 4096, needs d_a)
step8  1  K16         (writes slot 0 with d_b)
step9  0  K17         d_b → d_b   d_b → K17 ✓

The wrong-descale region is exactly the 1-2 iterations spanning the boundary. The wrong descale flows through m.thread_buf_[0] *= saved_k_descale[si] (L918) AND scale_s_k = scale_s * saved_k_descale[si] (L928), producing silent multiplicative-factor corruption in the attention output.

Note on the "first page index EVEN" filter — empirically, only even-first-page configurations push the per-element error above the FP8 threshold; odd-first-page configurations stay under threshold despite the same hazard mechanism (e.g. [4097, 4098] PASSes at max=0.0312, well above the all-page-4096 baseline of 0.0054 but below the 0.055 threshold). The exact reason — likely related to how the wrong descale value happens to combine with the running m_old and rowsum at the specific boundary iteration — is not fully derived here. Treat the parity condition as an empirical filter that selects the most visible failures, not as a load-bearing part of the root cause.

Cross-ROCm: bug reproduces on both ROCm 7.1 and 7.2 → source-level bug, not compiler/codegen.

Suggested fix (verified, but please feel free to take a different approach)

Below is the smallest fix I could find that resolves the hazard without changing the kernel's overall structure or losing the kFoldKDescale perf benefit. You know this kernel far better than I do — there may well be a cleaner approach (e.g. extending saved_*_descale to a 3-slot ring, separating K and V descale buffers, restructuring the iteration body, or something else I'm not seeing). Treat this as a starting point that demonstrates the bug is fixable with a localized change, not as a prescribed patch.

The idea is: defer save_descales to AFTER fmha_alu0 reads the slot, so the slot's previous content (descale of the K-tile being consumed) is preserved long enough. To get the right page index for the deferred save, capture k_physical_pages[number<0>{}] BEFORE K_page_issue advances it.

Properties of this approach:

• Pre-stage's three save_descales(0/1/0) calls are unchanged — pre-stage is manually unrolled without the pipeline lag, so it has no RAW hazard.
• kFoldKDescale performance benefit (no full-tile pk_mul_f32) is preserved.
• Deferred save can overlap with the next iteration's K_page_issue / cl_load, so the global-memory load latency is still hidden.
• ISA dump shows no register spill; SGPR/VGPR counts within 2 of baseline.

Diff (5 hunks, +20/-2 LOC):

diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_v3_pipeline.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_v3_pipeline.hpp
--- a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_v3_pipeline.hpp
+++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_v3_pipeline.hpp
@@ -875,6 +875,21 @@
             }
         };

+        // Deferred-save variant of save_descales. Used by the core_loop to
+        // defer the descale write until AFTER the same iteration's fmha_alu0
+        // has read the slot — closing a read-after-write hazard with the
+        // descale of the K-tile being CONSUMED this iteration (which lives
+        // in slot 1-pi from a previous iteration's save).
+        auto save_descales_at = [&](index_t captured_page, auto buf_idx_tag) {
+            if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::KV_BLOCKSCALE)
+            {
+                constexpr index_t b = decltype(buf_idx_tag)::value;
+                const index_t scale_offset =
+                    captured_page * nblock_stride_kv_block_descale +
+                    block_indices.kv_head_idx * nhead_stride_kv_block_descale;
+                saved_k_descale[b] = k_descale_ptr[scale_offset];
+                saved_v_descale[b] = v_descale_ptr[scale_offset];
+            }
+        };
+
         auto fmha_logits_trans = [&](auto sp_reg_idx) {
             if constexpr(kHasLogitsSoftCap)
             {
@@ -1270,7 +1285,9 @@
                     __builtin_amdgcn_sched_barrier(0);
                     __builtin_amdgcn_s_barrier();
                     __builtin_amdgcn_sched_barrier(0);
-                    save_descales(K_w0_lds_wr_idx);
+                    // Capture page for deferred save below; do NOT save here
+                    // (would overwrite the descale that fmha_alu0 needs to read).
+                    const index_t captured_page_w0 = k_physical_pages[number<0>{}];
                     K_page_issue();                                  // global_load_dword FIRST
                     __builtin_amdgcn_sched_barrier(0);               // prevent reorder
                     cl_load(memK, K_w0_lds_wr_idx, V_w0_lds_rd_idx); // buffer_loads SECOND
@@ -1296,6 +1313,8 @@
                     Scheduler::schedule(cl_p, number<2>{});
                     __builtin_amdgcn_sched_barrier(0);
                     fmha_alu_D_upd_pack();
+                    // Deferred save: fmha_alu0 has already read the slot.
+                    save_descales_at(captured_page_w0, K_w0_lds_wr_idx);

                     __builtin_amdgcn_sched_barrier(0);
                     // phase3
@@ -1367,7 +1386,9 @@
                     CK_TILE_FMHA_V3_ASM_MARKER("phase2 Wave4-7");
                     __builtin_amdgcn_s_barrier();
                     __builtin_amdgcn_sched_barrier(0);
-                    save_descales(K_w4_lds_wr_idx);
+                    // Capture page for deferred save below; do NOT save here
+                    // (would overwrite the descale that fmha_alu0 needs to read).
+                    const index_t captured_page_w4 = k_physical_pages[number<0>{}];
                     K_page_issue();                                  // global_load_dword FIRST
                     __builtin_amdgcn_sched_barrier(0);               // prevent reorder
                     cl_load(memK, K_w4_lds_wr_idx, V_w4_lds_rd_idx); // buffer_loads SECOND
@@ -1405,6 +1426,8 @@
                     Scheduler::schedule(cl_p, number<3>{});
                     __builtin_amdgcn_sched_barrier(0);
                     fmha_alu_D_upd_pack();
+                    // Deferred save: fmha_alu0 has already read the slot.
+                    save_descales_at(captured_page_w4, K_w4_lds_wr_idx);
                 }
                 return result;
             };

Verification — all 8 page-permutation probes pass after fix

pages	baseline	with fix
[4096]	PASS 0.0061	PASS 0.0061
[4097]	PASS 0.0063	PASS 0.0063
[4096, 4096]	PASS 0.0054	PASS 0.0054
[4096, 4097]	FAIL 0.0801	PASS 0.0084 ✅
[4098, 4099]	FAIL 0.1108	PASS 0.0088 ✅
[4097, 4098]	PASS 0.0312	PASS 0.0151
[4096, 4098]	PASS 0.0157	PASS 0.0054
[4095, 4096]	PASS 0.0247	PASS 0.0072

Note that the 3 already-passing multi-page cases also see significant max_diff reduction (3-4× lower), which independently confirms the fix attribution — all of them were silently affected by the same hazard, just below the FP8 threshold.

test_batch_prefill_4gb_boundary_targeted from aiter PR #3032 is a deterministic regression guard for this class of bug.