[HIP] [JIT] fp8_mqa_logits: hand-written gfx950 prefill indexer kernel

[anhcvt]

Motivation

The DeepSeek-V3.2 / GLM-5 "lightning indexer" produces the sparse-attention
selection logits. For each query row m and KV position n:

logits[m, n] = sum_h relu(<Q[m, h, :], K[n, :]>) * kv_scale[n] * weights[m, h]
               for n in [cu_seqlen_ks[m], cu_seqlen_ke[m]), -inf elsewhere

In the prefill path the caller has already gathered K out of the paged cache into
a contiguous [N, 128] buffer with a separate [N] fp32 scale, so there is no
block table and no preshuffle. This PR adds a hand-written HIP kernel for that
path on gfx950 (CDNA4), alongside the existing Triton/Gluon
aiter.ops.triton.attention.fp8_mqa_logits. The paged decode half is a separate
PR.

cu_seqlen_ks is a request's base offset into the gathered K buffer and
cu_seqlen_ke grows by one per query row, so each request is a causal triangle
and several requests arrive packed into one [M, N] chunk. Chunks are sized so
M*N*4 stays under VLLM_SPARSE_INDEXER_MAX_LOGITS_MB, which makes a 32K
request (4096, 32768) and a 128K one (1024, 131072).

Technical Details

New op: aiter/ops/fp8_mqa_logits.py

fp8_mqa_logits is a drop-in for the Triton entry point -- same tensors
(q_fp8, k_fp8, kv_scale, weights, cu_seqlen_ks, cu_seqlen_ke) and the
same clean_logits semantics. Key design points:

• No LDS staging. K is streamed HBM->register and contracted 32 heads x 32
  tokens per mfma_scale_f32_32x32x64_f8f6f4 tile.
• BLOCK_M query rows share one K stream, so K is read once per row block
  rather than once per row. Grid is (ceil(M / BLOCK_M), SplitN); SplitN
  splits each block's KV tile range so a small row grid still fills the device.
• Warps split the tile range, not the rows. That keeps several independent K
  streams in flight per block. Splitting rows across warps instead -- so K is read
  once per BLOCK_M * NUM_WARPS rows rather than once per BLOCK_M -- measured
  1.5-23% slower: the latency hiding is worth more than the L2 traffic it saves.
• Paired-row v_permlane32_swap_b32 head reduce. v_permlane32_swap_b32 a, b
  leaves a's reduction in lanes 0-31 and b's in lanes 32-63, so two query rows
  reduce with one swap and one add, and the store that follows uses all 64 lanes
  instead of half.
• kv-scale hoist. kv_scale >= 0 and ReLU is positive-homogeneous, so the
  scale is applied once per KV column after the head reduction rather than inside
  it.
• Row groups dispatched high-m first for N > 2048. Under causal masking a
  row group's work grows with m, so in natural order the longest-running blocks
  are dispatched LAST and become the tail; reversing starts them first and lets
  the short ones fill in behind. Worth +6% at (4096, 4096) and +3% at
  (8192, 8192), neutral elsewhere, and free -- it only reorders block dispatch.
• Fused -inf fill. With clean_logits the kernel writes the -inf outside
  each row's window itself. The Triton path pre-fills all M*N elements with
  torch.full and then overwrites the valid ones, paying for the valid region
  twice.
• -fno-honor-nans for the module, so the ReLU is a single v_max_f32.
  Without it LLVM must assume a signalling NaN and emits an IEEE canonicalize
  first -- two VALU per accumulator value, which is ~27% of the kernel's VALU.

BlockM, SplitN, num_warps, unroll2 and reverse_rows are all tunable;
zero means "use the host heuristic".

The kernel is gfx950-only and fixed at n_heads=32, head_dim=128 -- the shipped
GLM-5-FP8 indexer shape. is_supported(num_heads, head_dim) gates on that, so a
caller that also serves other shapes can route them to the Triton kernel rather
than trip a TORCH_CHECK. This mirrors how _should_use_asm_kernel gates the
head_size=128-only ASM paged-attention kernel in aiter/ops/attention.py.

Files added / changed:

• aiter/ops/fp8_mqa_logits.py -- the op and its support gate
• csrc/kernels/fp8_mqa_logits.cu -- kernel and host dispatch
• csrc/include/fp8_mqa_logits.h, csrc/pybind/fp8_mqa_logits_pybind.cu
• csrc/include/rocm_ops.hpp, aiter/jit/optCompilerConfig.json -- module_fp8_mqa_logits
• op_tests/test_fp8_mqa_logits.py -- correctness + perf sweep

Test Plan

op_tests/test_fp8_mqa_logits.py runs Triton and HIP on identical inputs and
grades both against one fp32 torch reference -- the same ref_fp8_mqa_logits the
Triton lane's test uses. Gates are an exact -inf mask match plus
calc_diff < 1e-3 and checkAllclose; tolerances are not widened.

The sweep is the cartesian product of 14 (s_q, s_k) shapes,
num_heads in {32,64,128}, head_dim in {64,128}, clean_logits in {0,1} and
six window modes -- 900 cases, 150 of which the HIP kernel supports. Points worth
calling out:

• Six window modes. Beyond causal and cp, the sweep covers misaligned,
  empty (rows with cu_ends below zero or below cu_starts), past_end
  (bounds beyond seq_len_kv) and multi_req (several requests packed into one
  chunk, so cu_starts jumps at each boundary and a block's rows straddle it).
  All are legal indexer input at a chunk boundary, and all are where the masking
  and the -inf fill are easiest to get wrong. Grading the HIP kernel under
  past_end caught two out-of-bounds writes during development, both on the
  clean_logits=False path.
• NaN-poisoned output. The output buffer is filled with NaN before each call,
  so a position the kernel fails to write fails the -inf mask check. Without it
  the check is close to vacuous -- the caching allocator hands back a block a
  previous case already left holding the correct -inf.
• The real chunk shapes: (4096, 32768), (2048, 65536), (1024, 131072).
  The reference runs in query-row chunks so its [heads, s_q, s_k] score tensor
  stays bounded (unchunked it is 17 GiB at heads=32, s_q=1024, s_k=131072).
• Cases the HIP kernel does not support leave its columns nan rather than
  reporting a wrong-but-fast number, and any case dropped for lack of memory is
  logged by name so a short table cannot read as full coverage.

python3 op_tests/test_fp8_mqa_logits.py

Test Result

All correctness gates pass on gfx950 across the sweep; per-case hip err matches
triton err. Grading the HIP kernel under past_end caught two out-of-bounds
writes on the clean_logits=False path (a cu_starts past seq_len_kv left the
fill's first range unclamped, and the store bounded abs_pos only by cu_ends);
both are fixed here.

Performance on MI355x/gfx950, num_heads=32, head_dim=128, run_perftest on an
otherwise idle GPU. causal is one request per chunk, multi_req is four:

s_q	s_k	clean_logits	window	Triton µs	HIP µs	speedup
1024	1024	True	causal	15.7	11.8	1.33x
4096	4096	True	causal	88.1	68.5	1.29x
8192	8192	True	causal	216.6	223.5	0.97x
4096	32768	True	causal	1063.6	613.4	1.73x
2048	65536	True	causal	1039.5	632.6	1.64x
1024	131072	True	causal	1608.2	820.6	1.96x
128	32768	True	causal	225.1	35.2	6.39x
671	131072	True	causal	1791.9	491.5	3.65x
4096	32768	True	multi_req	349.2	266.6	1.31x
2048	65536	True	multi_req	332.1	278.7	1.19x
1024	131072	True	multi_req	510.1	294.3	1.73x
128	32768	True	multi_req	63.5	21.6	2.94x
671	131072	True	multi_req	410.7	191.4	2.15x
8192	8192	True	multi_req	95.2	114.6	0.83x

Summarised over the 32-shape sweep (both clean_logits settings, both windows):

group	cases	geomean	range	HIP faster
causal	16	1.79x	0.87x - 6.39x	13 / 16
multi_req	16	1.31x	0.67x - 3.23x	11 / 16
clean_logits=True	16	1.64x	-	13 / 16
clean_logits=False	16	1.43x	-	11 / 16
all	32	1.53x	0.67x - 6.39x	24 / 32

The win tracks s_k / s_q, which is what the shared-K-stream design predicts:
BLOCK_M rows amortise one K read, so the longer a row's KV range is relative to
the row grid, the more there is to amortise. Every shape with s_k >= 8 * s_q is
a win (1.19x - 6.39x), and the largest is the small-M/long-KV corner
(128, 32768) at 6.4x, where the row grid alone cannot fill the device and
SplitN does the work.

The losses are the square, short-KV shapes -- (8192, 8192) at 0.97x/0.83x and
(1024, 1024)/(4096, 4096) under clean_logits=False -- where the tile loop is
short relative to the per-block Q/weights prologue. That prologue is also what
pins the kernel at 2 waves/SIMD: BLOCK_M=4 needs 194 VGPR, of which Q and the
per-row weights are 128, and both scale with BLOCK_M, so KV reuse and register
pressure cannot be traded apart in this design (BLOCK_M=8 needs 256 VGPR and
spills 113). Staging K through LDS would decouple them; until then
is_supported() plus a shape check lets a caller keep Triton on that corner.

clean_logits=True is the better case for the HIP kernel (1.64x vs 1.43x), which
is the fused -inf fill showing up: the Triton path pays a torch.full over all
s_q * s_k elements before the kernel overwrites the valid ones.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 1 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

[github-actions]

⚠️ [ruff] <PIE790> _{reported by reviewdog 🐶}
Unnecessary ... literal

reviewdog suggestion error

GitHub comment range and suggestion line range must be same. L65-L65 v.s. L65-L66