[HIP] [JIT] Fused MiniMaxM3 QKNorm+RoPE+CacheInsert

[weitliao]

Introduces minimax_m3_qknorm_rope_cache_shuffle_insert, a single HIP kernel that replaces the three-kernel sequence (fused QK-norm/RoPE -> reshape_and_cache(asm_layout=True) -> index-cache scatter) used by the MiniMax-M3 sparse attention layer, writing K/V directly into the page-16 SHUFFLE layout consumed by pa_decode_gluon.

Includes correctness tests against a torch reference and, where vLLM is available, end-to-end equivalence tests against the existing unfused path.

Motivation

The current 3-kernel path adds launch overhead and extra memory traffic (K/V written to qkv, then read again for cache insert). Fusing into one kernel reduces both on a hot prefill/decode path.

Technical Details

• Input: packed qkv row per token: [q | k | v | index_q | index_k], head_dim=128
• Processing: Gemma RMSNorm + partial NeoX RoPE on Q/K/index_q/index_k; V passes through unchanged
• Outputs: q_out, index_q_out, and paged k_cache / v_cache / index_cache in page-16 SHUFFLE layout
• Also supports: FP8 KV/index cache, padded slots (slot_mapping=-1), skip_index_branch=True

Test Plan

pytest op_tests/test_minimax_m3_fused_qknorm_rope_cache_shuffle.py -v
pytest op_tests/test_minimax_m3_fused_shuffle_vs_unfused.py -v # requires vLLM

Test Result

• Platform: linux, Python 3.12.13, pytest 9.1.1
• Command: pytest op_tests/test_minimax_m3_fused_qknorm_rope_cache_shuffle.py -v
• Result: 14 passed in 20.78s
  | Test | Cases | Result |
  |---|---|---|
  | test_fused_insert (tokens × fp8_index × quantized) | 1/7/64 × {False,True} × {False,True} (12) | PASSED |
  | test_skip_index_branch_leaves_index_outputs_untouched | 1 | PASSED |
  | test_multi_kv_head_and_full_rotary | 1 | PASSED |

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 4813 --add-label <label>

[andyluo7]

We exercised this exact implementation at commit 266922c417ceacb478317f56cf2001af106532fe in a MiniMax-M3 AgentX integration on MI355X (gfx950) and can contribute the following validation evidence.

Correctness:

• A direct GPU smoke covered both the full-index and skip-index branches and observed zero maximum absolute error for every compared Q, K, V, and index output.
• A later full real-verification GSM8K run completed 1,319/1,319 requests with no HTTP errors: strict exact match 0.9696739954510993, flexible exact match 0.9689158453373768.
• That full run included this fused cache-insert path plus the separately repaired vLLM unified-attention graph metadata, so it is integration evidence rather than an isolated attribution.

Performance:

• Controlled AgentX C1 job 822 changed only the MiniMax-M3 target QK-norm/RoPE/KV-insert path relative to its control.
• 147/147 measured requests completed with zero request errors.
• Across 115 exact matched requests, median TPOT improved 12.9% and median end-to-end latency improved 9.8%; 111/115 requests improved, and every input-length bucket improved by 12.7-13.4% at the median.
• Aggregate throughput was 5,294.14 generated tok/s/chip and P90 TPOT was 3.82 ms.

This is the same fused op carried by our downstream experiment, so we do not plan to open a duplicate AITER PR. The current PR is conflicting with main; I am happy to help rebase it onto current AITER main and attach a compact gfx950 reproduction if that would help move it forward.

AI disclosure: OpenAI Codex assisted with organizing the existing test evidence and drafting this comment; the reported numbers come from preserved run artifacts.

[andyluo7]

I prepared a current-main rebase helper branch without opening a duplicate PR:

• branch: https://github.com/andyluo7/aiter/tree/rebase/minimaxm3-fused-cache-insert
• head: 08587c2c0
• base: AITER main at f05d69d13276825f4f79a6b00a9c4f9ff1f20d1e

The six added source/test files are byte-identical to this PR's 266922c versions. The helper keeps main's composable_kernel pointer (15e12dd7f) and has no unrelated submodule delta. The only manual source conflict was relocating the pybind macro in the current csrc/include/rocm_ops.hpp context.

Static checks on the rebased tree pass: Black 26.3.0, Ruff 0.16.0, Python byte-compilation, and git diff --check. The earlier MI355X direct smoke and full GSM8K evidence in my prior comment exercised these same kernel/test source blobs.

Please feel free to use the branch to update this PR. I am intentionally not opening a competing AITER PR.

AI disclosure: OpenAI Codex assisted with applying and checking the rebase; original authorship remains Wei-Ting Liao.

[valarLip]

dulicated with fused_qknorm_idxrqknorm?

[weitliao]

dulicated with fused_qknorm_idxrqknorm?
@valarLip #4362 added the SGLang unit-scale NHD contract (fp8_e4m3_unit, BF16 index, block_size=1). This PR targets the vLLM/AgentX SHUFFLE path: page-16 layout for pa_decode_gluon, global/per-head FP8 scales matching reshape_and_cache, and skip_index_branch. Same optimization idea, different runtime contract — not a substitute for either side.