Add inverse_rope_group_quant op for DeepSeek-V4 wo_a input

[yzhou103]

Motivation

Fuses the inverse RoPE on the V4 attention output with 1xN e8m0 FP8 group quant that produces the wo_a GEMM input, saving a round trip through memory. Ships a HIP kernel plus a Triton implementation used as the test reference.

Kept as its own translation unit and its own JIT module: the op sits on the attention output path and shares no kernel or helper with the QK-norm/RoPE input-path ops, so folding it into their 6k-line file would only make every edit there rebuild it.

transpose_scale stores the scale column-major, which is what the preshuffled-B blockscale GEMM reads; the plain gemm_a8w8_blockscale wants row-major instead, so the flag tracks whether B is preshuffled in general. Both directions were checked numerically against the two GEMMs. .

Technical Details

Test Plan

Verified on gfx950 across quant group sizes 32/64/128, both scale layouts, and token counts 1..8192, matching the Triton reference on output and on the exact e8m0 scale bytes, plus HIP-graph capture/replay equivalence over all four kernel tiers.

Test Result

Measured on MI355X (gfx950, 256 CU), idle GPU, run_perftest hip_us over 101 iters.
Shape is the V4 attention output: o[S, H, 512] bf16 → x_fp8[S, G, 4096] + scale[S, G, 32] e8m0,
rope_dim=64, quant_group_size=128.

Baseline is the unfused two-kernel path this op replaces — triton inverse RoPE
(_rope_cached_bwd, rotating the rope tail in place) followed by
dynamic_per_group_scaled_quant with an fp8_e8m0 scale buffer, i.e. the same
pair atom's _V4RoPE.inverse path runs today.

hg = (n_local_heads, n_local_groups) = (16, 2) — V4-Pro TP8

S	fused µs	fused TB/s	unfused µs	unfused TB/s	speedup
1	2.45	0.010	4.73	0.005	1.93x
8	2.70	0.073	5.32	0.037	1.97x
32	2.63	0.302	5.42	0.146	2.06x
128	3.19	0.995	5.57	0.569	1.75x
512	4.39	2.888	6.75	1.879	1.54x
1024	6.25	4.058	9.15	2.773	1.46x
2048	10.94	4.635	14.66	3.460	1.34x
4096	18.64	5.441	23.48	4.320	1.26x
8192	33.10	6.130	40.58	5.000	1.23x
16384	70.77	5.734	89.78	4.520	1.27x

Peak 6.13 TB/s at S=8192, or 77% of the 8 TB/s HBM roofline. Fusing wins
~2x while launch-bound at low token counts and settles near 1.2x once
bandwidth-bound, where the win is the saved bf16 round trip rather than the
saved launch.

hg = (8, 1) — V4-Flash TP8, single group

S	fused µs	fused TB/s	unfused µs	unfused TB/s	speedup
1	2.35	0.005	4.47	0.003	1.90x
8	2.60	0.038	5.31	0.019	2.04x
32	2.65	0.151	5.48	0.073	2.07x
128	2.81	0.566	5.55	0.287	1.97x
512	3.51	1.818	5.98	1.066	1.71x
1024	4.59	2.776	6.76	1.884	1.47x
2048	6.23	4.091	9.19	2.775	1.47x
4096	11.50	4.434	14.83	3.438	1.29x
8192	19.03	5.359	23.71	4.300	1.25x
16384	34.37	5.934	41.43	4.923	1.21x

Quant group size axis — group 32 regresses at large S

hg=(16,2), row-major scale:

S	group	fused µs	fused TB/s	unfused µs	unfused TB/s	speedup
2048	32	13.09	3.907	15.20	3.363	1.16x
2048	64	11.08	4.591	15.27	3.331	1.38x
2048	128	11.70	4.335	15.38	3.298	1.31x
8192	32	42.48	4.814	41.13	4.971	0.97x
8192	64	34.67	5.868	41.05	4.956	1.18x
8192	128	34.25	5.924	41.05	4.943	1.20x

At S=8192 with group_size=32 the fused kernel loses (42.48 µs vs 41.13 µs) —
the only point in the sweep where it does. The unfused baseline is flat across
group size, since its quant kernel reads the same bytes either way; the fused
kernel is not, because group 32 does 4x as many amax reductions per row as
group 128 and the cross-lane reduction path does not amortize them at high
token counts. The wo_a path uses group 128, so this does not affect the
shipped configuration. Correctness is unaffected at every group size.

Byte count

TB/s = bytes / measured µs, where bytes = read o, plus one cos and one sin
row per token (all heads of a token share the row), plus written fp8 data at
1 B/elem plus one e8m0 scale byte per group. Both columns use the fused op's
ideal byte count, so the unfused figure is an effective bandwidth over the
same logical work — it really moves more, round-tripping the bf16 rows between
its two kernels. Bandwidth is not meaningful below S≈128, where both paths are
launch-bound.

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: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 4428 --add-label <label>

[Copilot]

Pull request overview

This PR adds a new DeepSeek-V4 output-path operator that fuses inverse (GPT-J style) RoPE on the attention output with FP8 (e8m0 scale) per-group quantization to produce the wo_a GEMM input, along with a correctness/perf test and the JIT/pybind plumbing to build and expose the HIP kernel.

Changes:

• Added a new HIP kernel (inverse_rope_group_quant) + C++ interface and pybind module for fused inverse-RoPE + FP8 group quant.
• Added a Python op wrapper that allocates outputs (including optional column-major scale layout via strides) and calls the compiled extension.
• Added an op test that validates fused vs reference (and includes an unfused baseline + optional HIP graph capture/replay check), plus JIT module registration and public export.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
op_tests/test_inverse_rope_group_quant.py	New correctness + benchmark + HIP-graph test harness for the fused op and an unfused baseline.
csrc/pybind/inverse_rope_group_quant_pybind.cu	New pybind translation unit to expose the C++ entrypoint and set stream handling.
csrc/kernels/inverse_rope_group_quant.cu	New HIP kernel + host launcher/validation for fused inverse RoPE and FP8 group quant.
csrc/include/rocm_ops.hpp	Adds the INVERSE_ROPE_GROUP_QUANT_PYBIND binding macro.
csrc/include/inverse_rope_group_quant.h	Declares the C++ API for the new op.
aiter/ops/inverse_rope_group_quant.py	New Python wrapper and @compile_ops entry for the new JIT module.
aiter/jit/optCompilerConfig.json	Registers the new JIT module sources/includes/flags.
aiter/jit/core.py	Excludes the new module from CK exclude-module logic (JIT plumbing).
aiter/init.py	Exports the new op at the package top level.

Comments suppressed due to low confidence (1)

csrc/kernels/inverse_rope_group_quant.cu:300

• positions is indexed assuming unit-stride contiguous storage (positions[s]), but there is no contiguity/stride check for it. If a caller passes a strided view (e.g. slicing a larger positions buffer), the kernel will read the wrong positions and silently produce wrong output.

  CHECK_CONTIGUOUS(o);
  CHECK_CONTIGUOUS(x_fp8);
  CHECK_CONTIGUOUS(cos_cache);
  CHECK_CONTIGUOUS(sin_cache);

---

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