[HIP] [JIT] fp8_paged_mqa_logits: hand-written gfx950 decode 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]

In the decode path K and its per-token dequant scale are co-packed in a paged
cache and gathered through a block table. This PR adds a hand-written HIP kernel
for that path on gfx950 (CDNA4), alongside the existing Triton/Gluon
deepgemm_fp8_paged_mqa_logits. The prefill half is a separate PR.

The decode indexer streams the entire K cache for a handful of FLOPs per token,
so it is squarely memory-bound and the grid carries a next_n axis -- every one
of a batch's MTP rows re-reads the same K. That redundancy, and the cost of
routing K through LDS, is what this kernel targets.

Technical Details

New op: aiter/ops/fp8_paged_mqa_logits.py

fp8_paged_mqa_logits is a drop-in for deepgemm_fp8_paged_mqa_logits -- same
tensors (q_fp8, kv_cache_fp8, weights, context_lens, block_tables,
max_model_len), same contract that only the causal window
p <= context_lens[b] - next_n + n is written and the -inf outside it is the
caller's. Key design points:

• No LDS on the K path. K is read from the paged cache straight to registers
  and contracted 32 heads x 32 tokens per mfma_scale_f32_32x32x64_f8f6f4 tile.
• ROWS_PER_BLOCK rows share one K stream. R consecutive next_n rows are
  handled by one block, so the cache is read once per R rows instead of once per
  row; grid is (batch, SplitKV, ceil(next_n / R)). R trades that redundancy
  against occupancy, so the host picks it -- along with ChunkK, num_warps and
  SplitKV -- from the total KV footprint batch * max_model_len. Each is
  overridable.
• kv-scale hoist. kv_scale >= 0 and ReLU is positive-homogeneous, so the
  scale is hoisted out of the head sum and applied once per KV column. This also
  matches the order the torch reference applies it in.
• v_permlane32_swap_b32 head reduce, instead of __shfl_down(x, 32) which
  lowers to ds_bpermute_b32 -- an LDS round-trip plus an lgkmcnt wait.
• KVBlockSize is a compile-time constant, instantiated for 1 and 64. Each is
  tied to one cache layout, the pairing production already uses: 1 reads the plain
  co-packed cache, 64 the shuffle_weight(layout=(16,16)) preshuffled one.
• -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, in the hottest loop of the kernel.

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, kv_block_size) 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_paged_mqa_logits.py -- the op and its support gate
• csrc/kernels/fp8_paged_mqa_logits.cu -- kernel and host dispatch
• csrc/include/fp8_paged_mqa_logits.h, csrc/pybind/fp8_paged_mqa_logits_pybind.cu
• csrc/include/rocm_ops.hpp, aiter/jit/optCompilerConfig.json -- module_fp8_paged_mqa_logits
• op_tests/test_fp8_paged_mqa_logits.py -- correctness + perf sweep

Test Plan

op_tests/test_fp8_paged_mqa_logits.py runs Triton and HIP on identical inputs
and grades both against one fp32 torch reference (a port of vLLM's
fp8_paged_mqa_logits_torch). 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 batch in {1,4,16,64},
next_n in {1,2,4,6}, heads in {32,64}, head_dim=128,
kv_len in {1024, 8192, 32768, 131072}, KVBlockSize in {1,64} and
var_ratio in {0.0, 0.3} -- 512 cases, 256 of which the HIP kernel supports.
Points worth calling out:

• next_n up to 6. The host clamps ROWS_PER_BLOCK to next_n, so a sweep
  stopping at 2 can never reach the R=3 instantiation MTP decode actually runs on.
• Ragged context (var_ratio 0.3 draws each length from +/-30% of kv_len).
  A uniform batch gives every sequence the same tail tile and the same causal
  boundary, so a kernel deriving its bounds from one sequence -- or from
  max_model_len -- would pass.
• Shuffled block pool. Block tables are built from a shuffled pool, so a
  kernel that ignores the table and walks the cache linearly fails rather than
  passing by accident.
• 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_paged_mqa_logits.py

Test Result

All correctness gates pass on gfx950 across the sweep; per-case hip err matches
triton err exactly.

Performance on MI355x/gfx950, heads=32, head_dim=128, ragged context
(var_ratio=0.3), run_perftest on an otherwise idle GPU -- 72 shapes:

B	next_n	ctx len	KVBlockSize	Triton µs	HIP µs	speedup
1	2	32768	1	15.8	5.4	2.90x
1	6	131072	1	65.9	14.3	4.61x
4	2	32768	1	40.5	10.2	3.99x
4	6	131072	1	312.9	40.2	7.77x
16	2	32768	1	166.6	29.2	5.71x
16	6	32768	1	394.8	46.9	8.43x
16	6	131072	1	1452.5	212.7	6.83x
64	2	131072	1	1987.8	594.0	3.35x
64	6	131072	1	5122.1	982.8	5.21x
1	2	32768	64	3.9	4.3	0.90x
4	6	131072	64	33.5	27.6	1.21x
16	6	32768	64	39.8	32.7	1.22x
16	6	131072	64	171.4	141.2	1.21x
64	6	8192	64	42.2	32.5	1.30x
64	6	131072	64	712.5	551.7	1.29x

Summarised over the full 72-shape sweep:

group	cases	geomean	range	HIP faster
KVBlockSize=1	36	3.86x	2.01x - 8.43x	36 / 36
KVBlockSize=64 (preshuffled)	36	1.01x	0.69x - 1.61x	18 / 36
all	72	1.97x	0.69x - 8.43x	54 / 72

The win is concentrated on the KVBlockSize=1 path, where the HIP kernel is
faster on every shape measured and the gap widens with next_n (geomean 1.60x at
next_n=1, 1.98x at 2, 2.44x at 6) -- which is what the shared-K-stream design
predicts, since R only has rows to amortise over once next_n > 1.

On the preshuffled KVBlockSize=64 path the two are at parity overall
(geomean 1.01x): the HIP kernel wins by ~1.2-1.3x at high next_n and loses by up
to 0.69x on the small end, where its prologue is not amortised. It is reported
here rather than hidden -- with is_supported() in place a caller can pick per
configuration, and the small-next_n preshuffled corner is the obvious next
target.

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 2 --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. L77-L77 v.s. L77-L78