[sumin-hong]

Validation status: Rebased onto ROCm/aiter main at f0b96f49e7d8518950292d41b6a27ceadc772315. The HIP kernel and launch behavior are unchanged from the implementation under review; Python formatting was adjusted for lint. The measurements below were supplied from their stated earlier revisions and have not been rerun on this rebased head. GPU correctness, performance and model evaluation on the new head remain pending.

[HIP] [JIT] fp8_mqa_logits: v35w gfx950 prefill indexer logits kernel

Motivation

This PR adds a faster HIP kernel for the prefill fp8 MQA indexer logits used by GLM-5.x / DeepSeek-V3.2 sparse attention, on gfx950 (MI355X):

logits[m, n] = sum_h relu(Q[m,h,:] . K[n,:]) * w[m,h] * kv_scale[n]
               for n in [ks[m], ke[m]), -inf elsewhere

The added API and JIT module are:

• Python entry: aiter.ops.fp8_mqa_logits.fp8_mqa_logits / is_supported
• JIT module: module_fp8_mqa_logits

It is an alternative to the Triton/Gluon aiter.ops.triton.attention.fp8_mqa_logits. It supports n_heads = 32 and head_dim = 128 only; other shapes are routed via is_supported().

Design (csrc/kernels/fp8_mqa_logits.cu)

• Workgroup layout. One workgroup is 8 waves × 4 rows = 32 query rows. Q and w stay in VGPRs, and all waves walk the same K sequence, so K is read once per 32 rows.
• K staging. K is staged through LDS with a bank-conflict-free swizzle, and kv_scale is staged next to it. Raw buffer loads prefetch 4 stages ahead. The vmcnt waits are placed by hand, because vmcnt also counts the logits stores and compiler-derived waits drain them every stage.
• Epilogue. MFMA is the plain v_mfma_f32_32x32x64_f8f6f4. The fp32 epilogue is scalar (v_max relu, v_fma with w) and is interleaved under the MFMAs. On gfx950, scalar VALU co-issues with an independent MFMA, but packed v_pk_* does not. Accumulation uses even/odd chains in a fixed order.
• Reduction and stores. The v_permlane32_swap head reduction is inline asm, because the clang builtin miscompiles its 2nd result on ROCm 7.2.3. kv_scale is applied after the reduction. The tail of each row pair is deferred by one sub-tile so it fills the MFMA→VALU hazard window, which leaves no s_nop in the loop.
• Masking and -inf. Stages are split into clean stages (no compare/select) and masked boundary stages, and both are shared across grid.y = SplitN. With clean_logits the kernel writes -inf itself, so there is no torch.full pass.
• Grid choice. (SplitN, reverse_rows) comes from a makespan model fitted on MI355X.
• Build flags. -fno-honor-nans -fno-slp-vectorize -mllvm -disable-machine-sink:
  • -fno-slp-vectorize: otherwise the epilogue is re-packed into v_pk_fma.
  • -mllvm -disable-machine-sink: otherwise the last row's MFMA pair sinks past the loop exit.

The table below describes the earlier v35w kernel. The current source head is the subsequent v35wes variant, which moves the stage barrier earlier and adds a per-wave s_setprio staircase. Its source commit reports bit-exact output versus v35w and 1.2–2.8% lower kernel time across two runs; those measurements have not been independently repeated here.

Kernel time on the GLM-5.3-Flash table shapes (MI355X, µs, clean_logits=True)

Setup:

• Harness: cold timing with a 1 GiB L2/MALL flush, device time from torch.profiler, median of 3 rotated runs.
• Image: vllm/vllm-openai-rocm:v0.30.0.
• Baselines: aiter main b59277d (14/10). FlyDSL uses flydsl 0.3.4.1; Gluon uses Triton 3.7.1.

s_q	s_k	FlyDSL (#4538)	this PR	Gluon (#5216)	HIP (#5046)	this / FlyDSL
causal · 1 request per chunk
8192	8192	181.5	142.3	219.0	228.7	1.28×
4096	32768	629.7	423.3	758.0	648.6	1.49×
2048	65536	633.1	430.2	790.3	723.5	1.47×
1024	131072	637.8	430.2	793.8	798.7	1.48×
16384	32000 (kpool 4)	2165.4	1571.5	2809.4	2388.6	1.38×
multi_req · 4 requests per chunk
4096	32768	218.6	155.8	268.0	274.7	1.40×
2048	65536	219.1	165.9	269.8	291.2	1.32×
1024	131072	221.5	204.1	273.5	282.9	1.09×
8192	8192	94.2	88.5	96.9	120.4	1.06×

Geomean vs FlyDSL #4538: 1.32× with clean_logits=True and 1.39× with clean_logits=False. It is also faster than every stronger baseline we measured, per shape:

• FlyDSL with the [HIP] [CI] [JIT] Add aiter.mqa_logits with a tuned Triton/FlyDSL table #6185 tuned table
• FlyDSL with the best of its 30 variants per shape
• FlyDSL [FlyDSL] fp8_mqa_logits: drop the host row-pad copies #6001
• Gluon built with the AMD Triton (3.7.0+amd.rocm7.2.0.git89002410)

Against the best of those on each shape, the geomean is 1.26× (clean) and 1.30× (no-clean). The one exception is 8192×8192 multi_req, where this PR loses by 3% (clean) and 18% (no-clean, vs Gluon on the AMD Triton).

Test plan

• Op test. python3 op_tests/test_fp8_mqa_logits.py --num-heads 32 --head-dim 128: all 150 HIP cases pass. They cover 6 window modes (causal, cp, misaligned, empty, past_end, multi_req), clean_logits 0/1, and a NaN-poisoned output buffer. The -inf mask matches exactly and checkAllclose err = 0.
• Op-test perf vs the Triton/Gluon kernel in main on the real chunk shapes (4096×32768, 2048×65536, 1024×131072): 1.38–1.85× (causal), 1.35–1.45× (multi_req), 1.11–1.57× (cp).
• JIT build identity. The JIT-built module is bit-exact against the standalone build used for the measurements and the vLLM E2E run, on 15 shapes × clean 0/1.
• E2E. In the vLLM nightly (GLM-5.3-Flash TP4, 128k/1k), output was correct: needles 5/5, top-512 identical on ≥ 99.8 % of rows vs the aiter result. TTFT changed by −1.4 % at c4.

Known limitations

• Default grid choice for multi-request chunks. The SplitN model assumes one causal request per chunk, so for multi-request chunks the default can be far from the best grid. For example, 1024×131072 multi_req takes 204 µs by default versus 159 µs with an explicit SplitN=20. Callers that know the layout can pass SplitN; a layout-aware heuristic is a follow-up.
• Small shapes. At M ≤ 128, or at small square shapes such as 1024×1024, the 32-row workgroup is too coarse and the Triton/FlyDSL kernels are faster. Route those through is_supported() plus a size gate in the caller if needed.
• Kernel-time scaling follows the windows of the whole 32-row block. Interleaved empty / past-end rows therefore cost about as much as full rows (the empty and past_end op-test modes).

Related work

• #4538 adds a FlyDSL prefill implementation with direct and LDS-staged variants; this PR adds an independent fixed-shape HIP implementation. The historical comparison above includes FlyDSL.
• #5216 optimizes the Gluon prefill path.
• #6185 proposes an AITER dispatch and tuned table for Triton/FlyDSL. This PR adds an opt-in HIP operator and does not integrate it into that dispatcher.
• #5046 was an earlier HIP implementation closed after comparison with existing backends; the table above includes that prior kernel.

Current validation

• Python files: Black 26.3.0 and Ruff 0.16.0 checks pass; Python AST and JIT configuration checked against the source implementation.
• GPU tests and benchmarks were not rerun after rebase.

AI assistance: Claude Code was used in the source development; OpenAI Codex was used for the rebase, static validation, and preparation of this PR.

PR Self-Review

Reviewed revision

460c230bae21403b6e1a2a48e3525a3003343249

Behavior and intent

This adds an opt-in gfx950 HIP prefill MQA logits operator. It computes weighted ReLU dot products over a contiguous gathered K buffer and supports clean or caller-managed output masks; existing Triton/Gluon and FlyDSL paths remain available. The aim is improved performance for the supported GLM indexer shapes.

Validation and unverified scope

On the rebased head, Python source AST and JIT configuration were compared to the source, the HIP kernel was checked against the prior PR head, and Black 26.3.0, Ruff 0.16.0, Python syntax and git diff --check passed. The supplied MI355X/ROCm measurements and commands are detailed above and predate this rebase. Unverified: GPU correctness, benchmarks and model evaluation have not been run on the current head; remote CI awaits maintainer approval.

PR-specific question

The intended contract is gfx950, 32 heads, head dimension 128, FP8 Q/K, float32 scales and weights, and int32 row bounds. The caller must check is_supported and supply the documented dtypes and contiguous layouts; the operator does not automatically route unsupported inputs. The C++ entry checks shape, contiguity and one-byte Q/K storage, but one-byte storage alone does not prove the FP8 dtype.

Changed op tests

op_tests/test_fp8_mqa_logits.py compares the new operator and Triton to the reference with clean and unclean outputs, poisoned buffers, causal, CP, misaligned, empty, past-end and multi-request windows. This new HIP entry and its masking behavior are not covered by the existing backend-only tests.

[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 (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
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: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
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 7183 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[aiter-bot]

NOT READY FOR HUMAN REVIEW

PR Self-Review

Revision 460c230bae21 · mode shadow

Author self-review accepted. Readiness report:

Author answers

• PASS — current-SHA answers are present and structurally specific

Show the author's self-review answers

Behavior and intent

This adds an opt-in gfx950 HIP prefill MQA logits operator. It computes weighted ReLU dot products over a contiguous gathered K buffer and supports clean or caller-managed output masks; existing Triton/Gluon and FlyDSL paths remain available. The aim is improved performance for the supported GLM indexer shapes.

Validation and unverified scope

On the rebased head, Python source AST and JIT configuration were compared to the source, the HIP kernel was checked against the prior PR head, and Black 26.3.0, Ruff 0.16.0, Python syntax and git diff --check passed. The supplied MI355X/ROCm measurements and commands are detailed above and predate this rebase. Unverified: GPU correctness, benchmarks and model evaluation have not been run on the current head; remote CI awaits maintainer approval.

PR-specific question

The intended contract is gfx950, 32 heads, head dimension 128, FP8 Q/K, float32 scales and weights, and int32 row bounds. The caller must check is_supported and supply the documented dtypes and contiguous layouts; the operator does not automatically route unsupported inputs. The C++ entry checks shape, contiguity and one-byte Q/K storage, but one-byte storage alone does not prove the FP8 dtype.

Changed op tests

op_tests/test_fp8_mqa_logits.py compares the new operator and Triton to the reference with clean and unclean outputs, poisoned buffers, causal, CP, misaligned, empty, past-end and multi-request windows. This new HIP entry and its masking behavior are not covered by the existing backend-only tests.

PR state

• PASS — pull request is marked ready for review

Copilot findings

• PASS — no active Copilot finding lacks a substantive author reply
• ADVISORY — Copilot could not review the current head (for example, quota or service failure)

Code and op-test hygiene

• PASS — no committed artifact or conflict marker detected

Required CI

• WAIT — Black: missing
• WAIT — Ruff: missing
• WAIT — Repository dependency: missing
• WAIT — Aiter Test Gate: missing
• Required CI is incomplete; this PR is not ready for human review.

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.