[Bugfix][MLA] Gate gfx942 native qh64 fp8 decode to page_size=64 (fix page_size=1 GPU fault)

[MohitAMD]

Motivation

The native QH64 fp8 persistent MLA-decode kernel added in #3188
(hsa/gfx942/mla/mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps.co) GPU
memory-access-faults on gfx942 (MI300X) for the page_size=1 decode shape.

aiter/mla.py routes gfx942 + nhead==64 + fp8/fp8 + max_seqlen_q==1 to this
kernel. #3188 validated correctness only at page_size=64. Production vLLM MLA
serves with page_size=1 (block_size=1), which the kernel never exercised — it
does an out-of-bounds access there, killing the worker on the first decode forward.

Surfaced on GLM-5.1-FP8 (GlmMoeDsaForCausalLM, MLA + DSA) 1P1D/2P2D
disaggregated serving, where per-rank head grouping is gqa_ratio=64
(data_parallel_size=8, tensor_parallel_size=1).

Closes #4363.

Reproduction

op_tests/test_mla_qh64_gfx942_pagesize1.py (added here) triggers the exact shape on
a single MI300X. On stock aiter (#3188 present, e.g. v0.1.18):

[aiter] LoadKernel: _ZN5aiter39mla_a8w8_qh64_qseqlen1_gqaratio64_v3_psE
[test] mla_decode_fwd persistent batch=1  kv_len=512  ... finite=True
[test] mla_decode_fwd persistent batch=4  kv_len=1024 ... finite=True
[test] mla_decode_fwd persistent batch=16 kv_len=2048 ...
Memory access fault by GPU node-2 ... on address 0x7f5bd0ed9000. Reason: Unknown.
Aborted (exit 134)

Confirmed on 3 independent aiter builds carrying #3188 (e03fa6040+04427a5d8,
ac90d5c89, origin/main tip) and on multiple MI300X nodes; the .co blob is
byte-identical from #3188 through tip, so the OOB is in the kernel binary.

Fix

Restrict the gfx942 native-qh64 dispatch clause in mla_decode_fwd to the validated
page_size==64. Other page sizes (page_size=1) fall through to the pre-#3188 qh16
fold (mla_a8w8_qh16_qseqlen1_gqaratio16_ps), which runs without fault. gfx950 is
unchanged.

or (
    get_gfx() == "gfx942"
    and nhead == 64
    and q.dtype == dtypes.fp8
    and kv_buffer.dtype == dtypes.fp8
    and max_seqlen_q == 1
    and page_size == 64      # native qh64 only where #3188 validated it
)

With the fix the repro folds to the qh16 kernel and completes (exit 0).

Alternatives considered

1. Drop gfx942 from the clause entirely — simplest, but discards Add native MLA QH64 fp8 persistent decode kernel for gfx942 #3188's benefit
   for the validated page_size=64 path.
2. Fix the kernel's page_size=1 OOB (preferred long-term) — the real root cause is
   an addressing/bounds bug in mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps at
   page_size=1. This dispatch guard is a safe stopgap; happy to re-enable the native
   path for page_size=1 once a fixed .co lands.

Test plan

• op_tests/test_mla_qh64_gfx942_pagesize1.py: aborts on stock Add native MLA QH64 fp8 persistent decode kernel for gfx942 #3188, passes with
  this fix (MI300X / gfx942).
• op_tests/test_mla_persistent.py -n 64,1 -d fp8 -kvd fp8 -blk 64 -b 1 -c 512 —
  no regression on the retained native page_size==64 path (MI300X / gfx942). Since
  this PR only removes page_size=1 from the gfx942 clause, page_size==64 dispatch
  is unchanged; verified on stock aiter (identical dispatch) that it loads the native
  mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps kernel, runs without fault, and passes
  the fp8 golden reference:
  [aiter] LoadKernel: _ZN5aiter39mla_a8w8_qh64_qseqlen1_gqaratio64_v3_psE ... .../gfx942/mla/mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps.co mla_decode-absorb_fp8 [golden fp8 vs aiter_asm] ... passed~ mla_decode-absorb_fp8 [golden fp8 split_out_ref vs aiter_asm] ... passed~
  i.e. the same kernel that OOB-faults at page_size=1 runs correctly at
  page_size=64 — the fault is page-size-specific, exactly as this guard assumes.
• Folded page_size=1 path validated end-to-end: GLM-5.1-FP8 DSA NIAH long-context
  retrieval to 96k on 1P1D EP8 and 2P2D EP16 (MI300X) with the fold in place — no
  fault, high retrieval accuracy.

Notes

This is the opposite head-count regime from #2821 (missing gqa=8 kernels at TP=8):
here the gqa=64 kernel exists but is broken at page_size=1.

[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: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 4365 --add-label <label>

[Copilot]

Pull request overview

This PR prevents a known gfx942 (MI300X) GPU memory-access fault by tightening MLA decode dispatch so the native QH64 FP8 persistent decode kernel is only selected for the validated page_size==64, while page_size==1 falls back to the pre-existing folded QH16 path.

Changes:

• Update aiter/mla.py dispatch logic to gate gfx942 native QH64 FP8 persistent decode to page_size==64 (gfx950 behavior unchanged).
• Add a targeted regression repro under op_tests/ that exercises the failing gfx942 + nhead=64 + fp8/fp8 + qlen=1 + page_size=1 decode shape.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
aiter/mla.py	Tightens gfx942 dispatch to avoid selecting the known-faulting native QH64 FP8 persistent decode kernel at page_size=1.
op_tests/test_mla_qh64_gfx942_pagesize1.py	Adds a standalone repro/regression script for the gfx942 page_size=1 faulting shape.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.