[OPUS] Absorb module_rmsnorm_quant into the opus rmsnorm module by carlushuang · Pull Request #4080 · ROCm/aiter

carlushuang · 2026-07-04T13:24:30Z

Summary

Depends on #4059 — this branch is stacked on carhuang/rmsnorm-opus (the #4059 head). Until #4059 merges, the diff here also shows its commits; it will collapse to just this PR's changes once #4059 lands. Please review/merge #4059 first.

This PR makes the single torch-free opus module (module_rmsnorm, introduced in #4059) the sole backend for the four public quant entrypoints as well:

aiter.rmsnorm
aiter.add_rmsnorm
aiter.rmsnorm_quant
aiter.add_rmsnorm_quant

They previously routed to the CK/HIP module_rmsnorm_quant (the shared add_rmsnorm_quant_kernel). They now run entirely on opus, so module_rmsnorm covers everything the old module_rmsnorm and module_rmsnorm_quant delivered, from one fast-building, torch-free translation unit set.

What is covered

The ported add_rmsnorm_quant_opus kernel reproduces every feature of the reference add_rmsnorm_quant_kernel:

no-quant (rmsnorm / add_rmsnorm), and dynamic quant to int8 / fp8 / fp4x2 (MXFP4)
per-token and grouped (per_1x128, per_1x32) quantization
fused residual-add, gemma norm, strided rows
shuffle_scale scale layout, and the fp4 e8m0 block-scale (non-shuffle + shuffle)

Implementation notes

The torch-free vectorized-IO + scaled-conversion layer lives in a self-contained csrc/include/opus/rmsnorm_opus_io.hpp that depends only on opus/opus.hpp. It copies the coalesced load_vector_nbytes / store_vector device functions and builds the quant conversions from opus primitives (opus::med3, opus::fp32_to_fp8_packed, opus::cast<fp4_t>). aiter_opus_plus.h is left untouched — no shared header between the rmsnorm rewrite and the attention kernels.
add_rmsnorm_quant_opus mirrors the reference's memory layout (interleaved per-token / contiguous grouped, gmem OOB), reduction (real DPP wave-reduce + single LDS exchange), packed v_pk_mul_f32 normalize, and fast reciprocal quant scaling (exact host 1/qmax, v_rcp inverse — no full IEEE divides). The group-max reduce uses a real-DPP warp_swap_ butterfly (quad_perm + upd_dpp) so it emits v_mov_b32_dpp rather than ds_bpermute.
Kernels are split into per-feature / per-dtype translation units under csrc/kernels/rmsnorm/ (norm bf16/fp16/fp32 + entry, quant, arq int8/fp8/fp4 + entry) so the JIT compiles them in parallel.

gfx1250

All translation units compile clean under --offload-arch=gfx1250 (opus is arch-portable and wave-size agnostic). The old module_rmsnorm_quant (CK / rocprim / hipcub) did not build for gfx1250, so absorbing it into opus also unblocks that target.

Feature coverage (all pass, gfx950 + gfx942)

test	what it covers	result
`test_rmsnorm2d.py`	rms_norm fp16/bf16/fp32, rmsnorm2d_fwd, fwd_with_add, T5/gemma, n up to 65536	810/810 pass (parity vs torch/CK)
`test_rmsnorm2dFusedAddQuant.py` modes 1-8	no-quant, per-token/grouped int8/fp8/fp4x2, fused-add, smoothquant, shuffle	pass (0 hard fails; the "warning" lines are fp8/fp4 rounding-boundary softness the CK reference shows identically)

The opus path also lifts the old module_rmsnorm_quant hidden-dim cap of 8192 (the 2d API now runs to n=65536+).

Compile time (clean JIT build, gfx950)

module	build
opus `module_rmsnorm` (this PR: 2d API and all 4 quant fns, split into parallel TUs)	~4 s
old `module_rmsnorm_quant` (separate module)	32.9 s

Clean builds (deps included). The opus module hosts both surfaces in one torch-free set of translation units (no c10 / rocprim / hipcub / ATen) and builds ~8x faster than the old quant module alone. Net effect: one fewer module to build and maintain. (For context, #4059 already replaced the CK module_rmsnorm — ~225 s — with the opus build.)

Performance vs `module_rmsnorm_quant` — full sweep, both arches

Measured with the interleaved-graph method (both kernels alternated in one process under CUDA graphs, so clock drift cancels and dispatch overhead is out of the path). % = opus / reference bandwidth, >100 = opus faster. Swept N = every model hidden_dim × M ∈ {1, 2, 4, …, 4096}.

gfx950 (MI350): all three paths (per-token rmsnorm_quant, per-token add_rmsnorm_quant, grouped per_1x128) are 97–102% across the entire N×M grid — effectively bitwise parity.
gfx942 (MI308X): per-token and grouped are 100–108% (opus at parity or faster); the add path is 97–100% except a handful of small-n large-m cells (n ≤ 2560, m ≥ 2048) at 92–96% — a residual latency-bound effect on those small-model shapes (see root cause below).

Non-quant path (the hot path): what the engines actually call

vLLM, SGLang and ATOM route their per-layer RMSNorm through two aiter entrypoints (bf16/fp16), gated by VLLM_ROCM_USE_AITER_RMSNORM / SGLang's _use_aiter:

rmsnorm2d_fwd — the input norm (no residual)
rmsnorm2d_fwd_with_add — the post-attention / post-MLP residual add + norm (called twice per layer, every model)

Both now dispatch the fp16/bf16, non-T5, hidden ≤ 8192, 2-D case to add_rmsnorm_quant_opus in no-quant mode — the same kernel main already dispatched here via the HIP add_rmsnorm_quant_kernel, and a 0-ULP bit-exact port of it. So the hot path runs the tuned single-pass kernel measured in the sweep above (97–102% on gfx950). fp32 / T5 (model_sensitive) / hidden > 8192 fall back to the generic opus norm kernel (≤2 ulp).

The generic fallback's out-of-place fused add was also fixed: it previously staged inputs with two host-side .copy_() passes before an in-place kernel (~1.5–2× slower than CK); it now runs a true single out-of-place pass via a compile-time OOP template, with the in-place / no-add instantiation byte-identical to before. This replaces the earlier bit-exact "be" kernel + bucket table (added for non-power-of-2 hidden), which are now removed entirely — the arq kernel covers those sizes exactly.

Capability win: grouped quant at n = 5120 / 6144

The old module_rmsnorm_quant rejects grouped per_1x128 at n = 5120 (GLM-4.5/4.6, Qwen3-14B/32B) and n = 6144 (MiniMax) — its dispatch picks a thread tile (TDS=24) that does not divide the group size, so it TORCH_CHECKs out. opus routes those to a divisor-clean tile and runs them correctly (982–4718 GB/s). So opus is strictly more capable on exactly the models that use block quant.

`n` = model hidden_dim: which shapes matter

n in rmsnorm is exactly the model hidden size. From current model configs:

hidden_dim (n)	models
1024 / 2048	Qwen3-0.6B / Qwen3-1.7B, Qwen3-30B-A3B
2560	Qwen3-4B
4096	Qwen3-8B, Qwen3-235B-A22B, MiMo-7B, DeepSeek-V4-flash
5120	Qwen3-14B, Qwen3-32B, GLM-4.5, GLM-4.6
6144	MiniMax-Text-01
7168	DeepSeek-V3 / V3.1 / V4-pro, Kimi-K2, Step-3

Every model that uses grouped / block quant (per_1x128 fp8, per_1x32 MXFP4 — DeepSeek, Kimi, GLM, MiniMax, Step) has hidden_dim ≥ 4096, where opus is at parity. The small-model shapes where the gfx942 add path dips (n ≤ 2560) run bf16 or per-token quant, both at parity.

gfx942 add-path fix: bf16 store rounding

The opus bf16 stores (residual_out, no-quant output) originally used round-to-nearest-even. gfx942 has no hardware bf16 cvt, so RNE lowered to a ~6-op/element software sequence (+181 VALU on the i8 256x32 add tile), dragging the gfx942 add path to ~85%. The CK/HIP reference truncates; switching opus to truncate matches it exactly and restores parity (gfx950 has the hardware cvt so it was unaffected either way). All op_tests still pass with truncate.

fused_qk_rmsnorm: no regression + strided rmsnorm fix

test_fused_qk_norm.py passes on both arches (144/144 q+k allclose checks). The real fused_qk_rmsnorm op is a separate module (module_fused_qk_norm_rope_cache_quant_shuffle) and is untouched by this PR.

Its test's split reference calls plain rmsnorm on torch.split views (row-strided). The opus 2d kernel read rows contiguously, so those views were materialized with .contiguous() — roughly 2x slower than the old CK kernel, which read the stride directly. Fixed by threading an input row-stride through the opus 2d norm kernels (output/residual stay contiguous); strided rmsnorm is now within ~0–9% of contiguous on both arches, and contiguous perf is unchanged (in_s == hidden).

Why the small-n large-m add gap persists (root cause)

Investigated with PMC counters + ISA. The group-max reduce originally used ds_bpermute (LDS); it was rewritten to a real-DPP warp_swap_ butterfly (quad_perm + upd_dpp with an uninitialized old value, which the compiler lowers to v_mov_b32_dpp instead of ds_bpermute on CDNA3). That cut grouped LDS instructions ~33% (9.8e4 → 6.5e4) — but the small-n timing barely moved, so ds_bpermute was not the bottleneck.

PMC at these small-n shapes: opus executes fewer VALU and SALU than the reference, yet takes a few % more active cycles at 100% GPU-util — i.e. it is latency/scheduling-bound (dependency bubbles from the extra reduction pass + scattered per-group scale stores), not instruction- or memory-bound. The hand-tuned rocprim kernel overlaps those stalls better at tiny row sizes; at n ≥ 4096 there is enough memory traffic to hide them, hence parity there. These small-n shapes run bf16 / per-token quant in real models (block quant is all hidden_dim ≥ 4096), so the residual gap is off every real path, and manual software-pipelining is not worth the complexity.

Impacted API

No signature changes. The four functions keep their signatures; only the backend changes. module_rmsnorm_quant is left in place (now unbound) as a fallback; its physical removal is a follow-up (it shares a header with gated_rmsnorm_quant).

The CK module_rmsnorm cold JIT build is ~225s on gfx950 (issue ROCm#4055): its blob_gen_cmd emits 1360 translation units / 4972 CK template instantiations, most of them quant / model-sensitive variants the common float path never uses. This adds a self-contained opus implementation of the plain RMSNorm path (rms_norm and fused residual-add), built as a single torch-free / pybind-free ctypes TU that compiles in ~1.8s on gfx950 (8 kernel instantiations). Build-time design: - Device kernels in csrc/include/opus/rmsnorm_opus_kernel.hpp with a host/device pass split so opus.hpp is parsed once (device pass only). - extern "C" C ABI via aiter_ctypes_error.h; tensors cross as the POD aiter_tensor_t, so there is no torch or pybind11 in the C++ world. - bf16/fp16, fp32 accumulation, any hidden size. A new env var AITER_RMSNORM_BACKEND (ck default, or opus) routes the public entrypoints rms_norm, rms_norm_cu, fused_add_rms_norm_cu, rmsnorm2d_fwd and rmsnorm2d_fwd_with_add through opus for the plain bf16/fp16 non-quant, non model-sensitive, non-gemma path; everything else falls back to CK. op_tests/test_rmsnorm_opus.py checks parity vs torch and CK across the matrix; op_tests/bench_rmsnorm_compile.py benchmarks the cold build wall of both modules. The opus kernel is HBM-bandwidth-bound on realistic shapes (5.2-8.2 TB/s on MI355X), so it does not regress against CK, which is bounded by the same roofline.

…TC__ Pass tensors as raw pointers + dims through the ctypes boundary instead of aiter_tensor_t, so the C++ side includes only opus/hip_minimal.hpp and opus.hpp (no <hip/hip_runtime.h>, aiter_tensor.h or aiter_ctypes_error.h). With -D__HIPCC_RTC__ the module preprocesses to ~14k lines and its cold build drops from ~1.6s to ~0.2s on gfx950. Validation and pointer/stream extraction move to the Python wrappers, whose public signatures are unchanged.

The normalize pass previously re-read the input from global, so on large hidden opus moved ~1.5x the bytes and lost ~15-22% to CK's register-tiled single-pass kernel. Cache up to CACHE_V vectors per thread (still writing the residual sum back for fused-add) so the row is read once; overflow beyond the cache reloads. Head-to-head vs CK ck_tile on MI355X (bf16), opus/CK throughput: 4096x8192 1.14, 2048x8192 1.22, 8192x4096 1.19, 8192x8192 1.06, 8192x2048 0.99, 8192x1024 0.85. Outputs match CK to bf16 rounding (~8e-3).

Extend the opus module to cover the whole CK module_rmsnorm surface so AITER_RMSNORM_BACKEND=opus serves it end to end: - dynamic & smooth quant to int8 / fp8, per-row yscale (rowmax/qmax) - fused-add and save-unquant variants - T5 / model-sensitive normalization on the plain and quant paths Runtime flags (residual/xscale/unquant pointers, quant/T5 ints) keep the instantiation count at 16 device kernels (in{bf16,fp16} x out{same,int8,fp8} x width{8,1}) in one torch-free TU; full-parity cold build is ~0.67s vs CK ~225s. Route the four quant entrypoints + T5 through opus in rmsnorm.py (group_size / shuffle_scale / gemma_norm still fall back to CK/quant). Also replace the block reduction with a deterministic sequential-addressing LDS reduction: the prior cross-warp shuffle path was intermittently nondeterministic on gfx950. Block size is rounded to a power of two and sized to the vector work. Validated on MI355X vs a CK-derived numpy reference across {int8,fp8} x {dynamic,smooth} x {add} x {save-unquant} x {T5} x {bf16,fp16}: yscale exact, int8 within 1 level, fp8 within fp8 granularity, residual bit-exact.

…educe) One row per block left small hidden launch/occupancy-bound (8192x1024 was ~0.68x CK). Use a 2D block: blockDim.x = threads-per-row, blockDim.y = rows-per-block, with a per-row (segmented) LDS reduction. tpr targets ~2 vectors/thread, so large hidden stays 1 row/block (unchanged) while small hidden packs several rows. opus/CK bf16 rms_norm on MI355X after the change: 8192x1024 0.68 -> 1.04, 16384x2048 0.93 -> 1.06, 8192x2048 0.88 -> 0.97; large hidden unchanged (2048x8192 1.11, 8192x4096 1.08, 8192x8192 1.02). opus is now >= 0.97x CK across the tested matrix. No ABI change (geometry only); correctness deterministic across all axes.

Fold the fused residual add into the plain kernel via a residual pointer (in-place when out == in), so there are two device kernels (norm + quant) instead of three, 12 instantiations instead of 16, and ~180 fewer lines. Unify the two block reductions into one templated helper and trim comments to one line where possible. No behavior change: full parity, deterministic, opus/CK bf16 rms_norm >= 0.96x across the matrix; full-parity cold build ~0.63s (was ~0.67s) vs CK ~225s.

Provide rmsnorm2d_fwd_opus and rmsnorm2d_fwd_with_add_opus with the same signatures (incl. use_model_sensitive_rmsnorm) as their _ck counterparts, so all four CK entrypoints -- rmsnorm2d_fwd, rmsnorm2d_fwd_with_add, rmsnorm2d_fwd_with_dynamicquant, rmsnorm2d_fwd_with_add_dynamicquant -- have a matching *_opus. The public dispatchers now pick _opus vs _ck symmetrically with identical args (opus covers T5 via use_model_sensitive_rmsnorm).

…n T5 use_model_sensitive_rmsnorm (T5) rounds the intermediate x*inv_rms to the storage dtype before applying gamma (to match vLLM's value distribution). CK's fused-add also differs by mode: the residual output is always round(x+res), but the norm uses the fp32 sum in the default path and the rounded sum in T5. opus previously used the rounded sum for fused-add in both modes, so default fused-add differed from CK on ~20% of elements. Cache the norm-input in fp32 and select fp32 sum (default) vs rounded sum (T5); plain and T5 paths are unchanged. Now every mode matches the CK formula (plain / fused x default / T5). No perf change (opus/CK >= 0.97x), compile ~0.66s, still 12 instantiations.

Add a bit-exact kernel (rmsnorm2d_fwd_be_kernel<scalar_t,TN,RN>) that reproduces CK's square_sum summation order using opus primitives -- not ck_tile: TN threads per row own RN width-8 vectors; the sum-of-squares is CK's intra-thread order (paired for T5, sequential for default) + a within-warp butterfly xor shuffle + a cross-warp tree over TN/64 warps. Because everything downstream (rsqrtf, T5 dtype-rounding, gamma, residual, quant) already matched, reproducing square_sum makes the output bit-identical. launch_norm dispatches hidden in {64,128,512,1024,1536,2048,3072,4096,6144,8192} (CK's vn=8 buckets) to the matching (TN,RN); other sizes and quant keep the fast generic path (formula-exact, <=2 ulp). Verified on MI355X vs the real CK kernel: bf16 T5 and default are 100% bit-identical across all buckets (25M-100M elems, multiple seeds). fp16 is deterministic and within 2 ulp (99.995%); a residual 1-ulp square_sum difference that bf16 rounding absorbs but fp16 exposes -- closing it needs CK's exact warp butterfly derivative, tracked as follow-up. Compile ~0.95s / 32 instances.

Storing the norm-input in an ext_vector let the compiler reorder the squared-sum; a plain float[RN][8] keeps the summation order, halving fp16's deviation from CK (910 -> 205 elems of 16M, still <=2 ulp). bf16 stays 100% bit-identical.

fp16's residual ~1-ulp vs CK is TU-context codegen (identical source is bit-exact in isolation), not the intra-thread FMA form. bf16 stays 100% bit-identical; fp16 is <=2 ulp. Reverted to the clearer expression.

…ault Merge all opus Python bindings/wrappers into aiter/ops/rmsnorm.py and delete the separate rmsnorm_opus.py, so rmsnorm.py is the single place for the op. The opus wrappers (rms_norm_opus, fused_add_rms_norm_opus, rmsnorm2d_fwd_opus, rmsnorm2d_fwd_with_add_opus, *_dynamicquant_opus, *_smoothquant_opus) are a complete bf16/fp16 implementation covering plain / fused-add / dynamic+smooth quant (int8/fp8) / T5, any hidden size -- so the CK (_ck) functions are now a removable opt-in. AITER_RMSNORM_BACKEND now defaults to 'opus' (set =ck for the legacy CK path). With opus default and self-contained, all CK bindings can be deleted later with no functional loss (gemma_norm / group_size / shuffle_scale keep using the separate module_rmsnorm_quant, not CK). C++ TU (module_rmsnorm_opus) unchanged.

- add fp32 for norm and quant so opus covers every CK dtype (width 16/sizeof: 8 for bf16/fp16, 4 for fp32; bit-exact be-kernel stays 2-byte) - element<->fp32 via opus::cast (drop the to_f32/from_f32 shims); pin OPUS_FP32_to_BF16_DEFAULT=0 so bf16 rounds RNE on every arch - each kernel now takes a single Traits param (fwd/quant/be traits) - gfx942 fixes: barrier before reusing the block_reduce LDS buffer; cap threads-per-row at 256 - condense comments

opus now serves fp16/bf16/fp32 for the plain, fused-add, dynamic/smooth quant and T5 paths at any hidden size, so the CK rmsnorm is redundant. - delete the eight *_ck bindings and the module_rmsnorm dependency - drop the AITER_RMSNORM_BACKEND switch (opus is the only backend); gemma_norm/group_size/shuffle_scale still fall back to module_rmsnorm_quant - update op_tests to the non-ck entrypoints and a torch reference

Instead of adding a separate module_rmsnorm_opus alongside the (now dead) CK module_rmsnorm, point module_rmsnorm itself at the opus source so the existing module name builds the fast, torch-free ctypes TU. - optCompilerConfig: module_rmsnorm now builds rmsnorm_opus_kernels.cu (-D__HIPCC_RTC__, no blob-gen); drop the module_rmsnorm_opus entry - rmsnorm.py: opus @compile_ops target back to "module_rmsnorm" - delete the orphaned CK sources (rmsnorm_kernels.cu, rmsnorm_ck_kernels.cu, rmsnorm_pybind.cu, rmsnorm.h) and the RMSNORM_PYBIND macro - point the build-wall / compile bench at module_rmsnorm Verified end-to-end (gfx950, ROCm 7.2.2): rms_norm builds module_rmsnorm.so in 1.4s (was ~225s for the CK build), cached call 0.21ms.

Rename csrc/include/rmsnorm_opus.h back to the original rmsnorm.h so the opus impl reuses the existing header name instead of adding an _opus one.

The T5 case passed use_model_sensitive_rmsnorm as the 7th positional arg of rmsnorm2d_fwd_with_add, which is gemma_norm. With the opus backend that routes gemma_norm=True to module_rmsnorm_quant (add_rmsnorm), whose kernel only supports n<=8192 (TORCH_CHECK(false) otherwise) -> the n=16384/32768/65536 cases crashed. Passing the arg by keyword routes the call to the opus model-sensitive path, which handles any hidden size, so the test now exercises T5 as intended.

gemma_norm now runs on the opus norm kernel (weight+1) at any hidden size, so rmsnorm2d_fwd_with_add(gemma_norm=True) no longer falls back to the shared module_rmsnorm_quant kernel (which caps at n<=8192 and would TORCH_CHECK-crash). - opus kernel: gemma is a compile-time template param (if constexpr), so gemma=0 is byte-identical to the previous kernel (verified bit-exact + 1.000x perf on gfx950); only gemma=1 adds the +1 offset. BE bit-exact path is untouched (gemma uses the generic kernel, any n). - C ABI: rms_norm_opus / fused_add_rms_norm_opus gain an int gemma arg. - dispatch: _use_opus no longer excludes gemma; group_size/shuffle_scale (grouped/ MXFP4 quant, which legitimately live in the shared module_rmsnorm_quant) and exotic dtypes keep the fallback, now with an explicit hidden<=8192 assert instead of a cryptic kernel abort. - test_rmsnorm_opus.py: add a gemma_norm parity case (covers n>8192). Validated gfx950 + gfx942: gemma=1 matches rmsnorm*(weight+1) for n up to 65536; compile stays ~1.4s (single TU, 44 instances).

…template param Traits already holds every compile-time kernel parameter (scalar_t, width), so fold the gemma flag into fwd_traits<Scalar, Width, Gemma> and read Traits::gemma in the kernel. rmsnorm2d_fwd_kernel is back to a single 'typename Traits' param; launch_norm selects fwd_traits<..., true/false>. Pure refactor: same instances, gemma=0 still byte-identical to pre-gemma and gemma=1 unchanged (re-verified bit-exact + 1.000x perf on gfx950).

Flatten into namespace aiter and give the kernels/traits self-describing names: rmsnorm2d_fwd_kernel -> rmsnorm_opus_kernel rmsnorm2d_quant_kernel -> rmsnorm_quant_opus rmsnorm2d_fwd_be_kernel -> rmsnorm_be_opus fwd_traits/quant_traits/be_traits -> rmsnorm_opus_traits/rmsnorm_quant_opus_traits/rmsnorm_be_opus_traits Pure rename (no ABI/behavior change); the extern C entrypoints are unchanged. Re-verified gfx950: gemma=0 bit-identical, gemma=1 correct, dynamic-quant int8 correct, perf 1.000x; builds on gfx942.

Drop the launch_dims struct; pick_dims now returns std::pair<dim3,dim3> (block, grid) and callers use 'const auto [block, grid] = pick_dims(...)'. Pure refactor; hipLaunchKernelGGL is a direct <<<>>> macro (no lambda capture) so the bindings are fine under C++17. Re-verified gfx950 (gemma bit-identical + quant + 1.000x perf), builds gfx942.

AMDGPU tolerates misaligned 128-bit global access (verified bit-exact, no fault, down to 2-byte offset on gfx942 and gfx950 across the BE, generic and quant paths), and tensor pointers are always at least element-aligned. So the vec path is chosen purely on length (hidden % VW == 0); aligned16() and the per-pointer checks are removed. Aligned inputs are byte-identical and same perf; misaligned inputs now take the fast vec path instead of the scalar fallback.

The kernel impl works in template element types (Traits::scalar_t/in_t/out_t) plus builtin float, so it needs no element-type aliases. Replace opus::cast<fp32_t> with opus::cast<float> and i8_t with signed char in the kernel, and move the host-facing vocabulary (bf16_t/fp16_t/fp32_t/i8_t/fp8_t, used only to instantiate the launchers) into rmsnorm.h. Pure relocation: float==fp32_t, signed char==i8_t; re-verified bit-identical + quant + 1.000x perf on gfx950, builds gfx942.

ROCm#4056) ROCm#4056 gated the device-only TDM builtin behind the host pass too, so opus.hpp now compiles on the host pass. Include it unguarded and source the launcher's element vocabulary from opus (using bf16_t = opus::bf16_t; ...) instead of redefining the types with a duplicated clang-version #if. Single source of truth, no drift. Compile time is unchanged (gfx950 1.46s; gfx942 4.27->4.33s, +0.06s), correctness bit-identical (gemma=0), gemma=1/quant correct, perf 1.000x; builds clang-20+clang-22.

Replace the hardcoded 127/448/240 in _qmax_outcode with aiter.ops.quant.get_dtype_max (torch finfo/iinfo), keeping the int8/fp8 support guard and the out_code (0=int8, 1=fp8). Values verified identical (127/448/240).

The plain / fused-add / gemma entrypoints always use opus now: the old fallback (module_rmsnorm_quant) only supports fp16/bf16, a subset of opus's fp16/bf16/fp32, so it served no dtype opus doesn't already handle, and opus's own _check gives a clear error for unsupported dtypes. Removed _use_opus and the dead fallbacks; the dynamic-quant paths now gate purely on the real feature (group_size/shuffle_scale -> shared module_rmsnorm_quant, hidden<=8192). Routing verified via stub dispatch.

New csrc/include/opus/rmsnorm_opus_quant_detail.hpp includes only opus.hpp (no torch/rocprim/hipcub/aiter_opus_plus.h) so the module_rmsnorm TU stays a torch-free ~1.4s single-TU build. Inlines the MXFP4 e8m0 block-scale (ceil_pow2(amax/6)), e8m0 byte, and the mx_scale_shuffle_idx swizzle from mx_quant_utils.h. Verified on gfx950: builds at 1.47s (no regression) and both helpers are bit-exact vs reference. Foundation for absorbing module_rmsnorm_quant (grouped/shuffle/fp4) into opus.

ds_bpermute (shfl) has ~20x the latency of a DPP ALU cross-lane op, and on the per-token / no-add quant paths the block reduce sits on the critical path. Add a wave64 DPP all-reduce (bound_ctrl 0-fill is a valid identity: sum over squares, max over |values|, both non-negative) and use it in block_reduce_1d; keep the shfl butterfly as the wave32 fallback. gfx950 vs module_rmsnorm_quant now: n>=4096 shapes 86-99%, add path 90-99%, grouped 86-99%, fp4 90-97% of reference bandwidth.

Match the reference's packed-f32 compute for the square-sum and the inv/weight multiplies. Lifts grouped small-n (n<=2048) a few points; neutral on the memory-bound shapes. Correctness unchanged (exact).

Profiling small-n no-add per-token (rocprofv3, gfx950) showed identical memory traffic to module_rmsnorm_quant but ~4x the LDS instructions: the cross-warp part of block_reduce_1d used opus::shfl (ds_bpermute = LDS) for a full 6-step butterfly on the per-warp partials. Combine the (compile-time NWARP<=16) partials serially from LDS instead -> LDS insts now match the reference exactly. Also make the square-sum scalar (v_fmac) like the reference (packed pk_mul added ops); keep the packed normalize. gfx950 rocprofv3 vs reference (no-add per-token): n=2048 76%->91%, and n>=4096 / add paths at 96-100%.

Each <IS_MAX,BLK> instantiation owns a distinct __shared__ array, so the sum and max block-reduces can never race on it -- the leading sync_threads was unnecessary. One barrier per reduce now, matching the reference block_reduce.

…head ISA diff vs module_rmsnorm_quant (n=2048 no-add) showed opus carried full IEEE divides (v_div_scale_f32 x ~3) the reference avoids: it stores max*(1/qmax) and inverts via v_rcp. Replace all three divides with fast reciprocals: - scale ys = max * rcpf(qmax); inv_ys = rcpf(ys) - rms mean via sumsq * rcpf(n) (v_rcp_iflag, like the reference) and cache the per-device CU count in python (get_device_properties per call is costly for the few-us small-n kernels). v_div_scale now 0, kernel VALU below the reference. rocprofv3 gfx950 no-add per-token: n=2048 90%->~96%, n=8192 ~101%, others 97-99% of reference.

The previous commit used an in-kernel rcpf(qmax) for the stored quant scale; the approximate reciprocal flips a few int8 roundings at boundaries (data-dependent, caught on gfx942). Pass the EXACT 1/qmax from python instead (still no divide, bit-accurate scale, matching the reference's compile-time 1/qmax). Also actually land the per-device cu_num cache (lost in an earlier A/B restore). Re-verified: gfx950 8-seed exhaustive int8 PASS; gfx942 int8/fp8/no-quant PASS; timing unchanged (n=2048 ~97%, n>=4096 98-101% of module_rmsnorm_quant).

github-actions · 2026-07-04T13:24:59Z

🏷️ 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 4080 --add-label <label>

The grouped group-max used opus::shfl (ds_bpermute) for its butterfly. Replace it with the warp_swap_ butterfly (from carlushuang/gcnasm warp_sort_bitonic): quad_perm for lanegroup 2/4 and upd_dpp with an uninitialized old value + complementary bank masks for 8/16, which the compiler lowers to real v_mov_b32_dpp instead of falling back to ds_bpermute on CDNA3 (rts<=16 -> zero LDS group reduce). Cuts grouped LDS instructions ~33% (9.8e4 -> 6.5e4). Block reduce kept as-is: its row_shr/row_bcast already fuse into v_add_f32_dpp/v_mov_b32_dpp (real DPP); a shfl butterfly there would add ds_bpermute for the 32/64 span. Correctness exact (8-seed exhaustive).

rmsnorm2d_fwd fed a torch.split view (row stride != hidden, e.g. q/k from fused_qk_rmsnorm) hit 'rms_norm_opus: contiguous only'. The opus kernel reads rows contiguously, so materialize a non-contiguous input first. Fixes test_fused_qk_norm.py.

_arq asserted input.is_contiguous(), but the add_rmsnorm_quant_opus kernel already takes the input row stride (in_s). fused_qk_rmsnorm calls aiter.rmsnorm on a torch.split view (row-strided, last dim contiguous), which now works via the stride; only materialize when the last dim itself is strided.

…ale store Profiling (rocprofv3 GRBM + a disable-the-block experiment) traced the ~5% grouped overhead vs module_rmsnorm_quant to the per-group scale-store block. The store predicate used a runtime integer modulo (tid % rts) evaluated by ALL threads (no HW integer divide on GPU); rts and group_size are always powers of two, so use a mask and shifts instead. Recovers ~1%. The residual is the scattered global scale store's scheduling, which the hand-tuned reference overlaps better (grouped opus ~96% of reference, deterministic GRBM; per-token is at parity).

Two correctness bugs surfaced benchmarking real hidden dims (5120/6144/7168): 1. Grouped dispatch: for n in (4096,6144] the per-token tile (256,24) was used, whose thread_data_size=24 does NOT divide group_size (128/32) -> non-power-of-two rts -> ctz()=0 -> out-of-bounds scale store -> GPU fault (e.g. n=5120 = GLM hidden). Route grouped n>4096 to the (512,16)/(1024,8) tiles (TDS divides the group) for ALL such n, not just n>6144. This also *adds* support the reference TORCH_CHECK-rejected. 2. fp4 output OOB bound: oob_o used n, but an fp4 row is n/2 bytes. When n < BLK*TDS (n=5120/6144/7168) OOB threads stored past the fp4 output row into adjacent memory (latent memory corruption; only safe at exact n=4096/8192). Size the buffer by the real element count (n/2 for fp4). Verified gfx950: grouped int8 + fp4 (shuffle/non-shuffle) now pass for n in {4096,5120,6144,7168,8192} x grp {32,128}; per-token/no-quant unaffected; gfx942 builds.

… (no gfx942 regression) The residual_out/no-quant bf16 stores used round-to-nearest-even, which has no hardware cvt on gfx942 and lowers to a ~6-op/element software sequence (+181 VALU on the i8 256x32 add tile), making the add path ~85% of the CK/HIP reference there. The reference truncates; switching to truncate matches it exactly. gfx950 (hardware bf16 cvt) is unaffected. Add path is now 97-100% of old on both gfx942 and gfx950; all op_tests (rmsnorm_opus, rmsnorm2d, rmsnorm2dFusedAddQuant) still pass.

rms_norm_opus previously materialized any non-contiguous input via .contiguous() so the opus 2d kernel could read rows contiguously. For a row-strided view (e.g. a torch.split slice feeding fused_qk_rmsnorm) that copy roughly doubled the time vs the old CK kernel, which read the stride directly. Add an input row-stride (in_s) to the be/generic no-quant kernels and their launchers (output/residual stay contiguous), and pass input.stride(-2) for a 2-D row-contiguous view instead of copying. Strided rmsnorm is now within ~0-9% of contiguous (was ~2x); the real fused_qk_rmsnorm op is a separate module and was never affected. Contiguous perf unchanged (in_s==hidden).

…rch.compile The opus backend is ctypes and reads .data_ptr() in Python, so torch.compile traced into rms_norm_opus and hit 'Cannot access data pointer of FakeTensor' (broke ATOM gpt-oss / DeepSeek / Kimi serving, which compile the model). Wrap the public rmsnorm entrypoints (rms_norm, rmsnorm2d_fwd, rms_norm_cu, fused_add_rms_norm_cu, rmsnorm2d_fwd_with_add) with torch_compile_guard + a fake impl so they are opaque aiter custom ops, exactly as the pre-opus CK ops were registered. Eager unchanged.

… for 2560/5120/7168 rmsnorm2d_fwd_with_add (the per-layer residual-add that vLLM/SGLang/ATOM all call) was implemented as two host-side .copy_() staging passes feeding the in-place opus kernel -- ~1.5-2x slower than the CK reference on that hot path. Add a true out-of-place fused-add: the be/generic norm kernels gain a compile-time OOP template that reads input/residual_in and writes out/residual_out in a single pass (OOP=false keeps the in-place / no-add instantiation byte-identical, so those paths are unchanged). New add_rms_norm_opus C ABI entrypoint; rmsnorm2d_fwd_with_add_opus calls it directly (no copies), covering all hidden sizes and the T5 variant. Also add bit-exact be-kernel tiles for hidden 2560/5120/7168 (Qwen3-4B, GLM-4.5/4.6 & Qwen3-14B/32B, DeepSeek/Kimi/Step) so these non-power-of-2 sizes hit the tuned kernel instead of the generic one. Result on gfx950: no-add and fused-add are now at parity with CK (97-108%) across the full hidden_dim x M grid, both bf16/fp16; all op_tests pass.

Everything was in one translation unit (csrc/py_itfs_cu/rmsnorm_opus_kernels.cu, ~328 kernels), so the whole module compiled serially. Move the kernels into csrc/kernels/rmsnorm/ split by feature so ninja builds them in parallel: rmsnorm_opus_norm.cu - rms_norm / fused_add / add_rms_norm (launch_norm) rmsnorm_opus_quant.cu - rms_norm_quant (dynamic/smooth int8/fp8, launch_quant) rmsnorm_opus_arq_i8.cu / _fp8.cu / _fp4.cu - add_rmsnorm_quant per out dtype rmsnorm_opus_arq_entry.cu - the C entrypoint + no-quant + out_code dispatch The out_code dispatch that lived in launch_arq is replaced by per-out-dtype launchers (opus_arq_i8/fp8/fp4/noquant) so int8/fp8/fp4 instantiate in separate TUs. The kernel templates (rmsnorm_opus_kernel.hpp) and launchers (launch_norm/launch_quant/launch_arq_io) are unchanged, so the exact same kernels are selected by the same runtime dispatch -- no functional or perf change. Clean parallel compile ~2s vs ~5s single-TU on gfx950; all op_tests pass; all 6 TUs compile for gfx1250.

module_rmsnorm was one translation unit (csrc/py_itfs_cu/rmsnorm_opus_kernels.cu). Split the kernels into csrc/kernels/rmsnorm/ by feature so ninja builds them in parallel: rmsnorm_opus_norm.cu - rms_norm / fused_add / add_rms_norm (launch_norm) rmsnorm_opus_quant.cu - rms_norm_quant (dynamic/smooth int8/fp8, launch_quant) The kernel templates and launchers are unchanged, so the same kernels are selected by the same dispatch -- no functional or perf change; op_tests pass, torch.compile works. (ROCm#4080 adds the add_rmsnorm_quant arq units on top; keeping the same layout here.)

norm was the compile bottleneck (~2.6s in one TU): its be+generic kernels are the expensive ones to instantiate. Split launch_norm by input dtype into separate TUs via the same extern-dispatch pattern as arq -- opus_norm_bf16/fp16/fp32 each in their own .cu, with the entrypoints (rmsnorm_opus_norm_entry.cu) dispatching the dtype code to them. Kernels/launcher unchanged -> identical kernels, no functional/perf change. Parallel compile now ~1.65s (was 2.65s after the first split, 4.58s single-TU) on gfx950 -- ~2.8x off the original. New bottleneck is norm_bf16/fp16 (~1.6s, mostly the 0.76s header parse). op_tests pass, torch.compile works, all TUs compile for gfx1250.

The quant rewrite pulled its coalesced load/store + scaled conversions from a copy of aiter_opus_plus.h's layer (opus_vec_io.hpp). Instead: revert aiter_opus_plus.h to untouched, and give the rewrite its own small self-contained header (rmsnorm_opus_io.hpp, opus.hpp-only) that copies just the generic load/store_vector device functions and does the output conversion via opus internal API -- opus::cast (bf16/fp16), opus::med3 + opus::fp32_to_fp8_packed (fp8), opus::fp32_to_i8 (int8), opus::cast<fp4_t> / fp32_to_fp4_packed (fp4). Drops the scaled_cast / fp32_to_*_scaled / bf8 helper layer. Same instructions (med3+cvt_pk_fp8 etc.) -> identical numerics/perf; op_tests pass, quant paths at parity, gfx1250 compiles.

…old coverage into test_rmsnorm2d The refactor routes all rmsnorm/quant public APIs through opus, so the existing test_rmsnorm2d.py and test_rmsnorm2dFusedAddQuant.py already exercise the opus backend. Remove the extra opus-specific test and the compile-time bench (compile time is measured locally). Fold the two cases test_rmsnorm2d didn't have -- fp32 input and gemma_norm -- into test_rmsnorm2d so no coverage is lost. All checks pass (810 in test_rmsnorm2d; int8/fp8/fp4 per-token/grouped/shuffle in FusedAddQuant).

The bf16/fp16 non-T5 hidden<=8192 case that main dispatched to the HIP add_rmsnorm_quant_kernel now goes through the bit-exact opus arq kernel (add_rmsnorm_quant_opus, 0 ULP vs the HIP kernel) instead of a dedicated bit-exact 'be' kernel. The generic kernel (<=2 ulp) handles the fallback (fp32, T5/model_sensitive, hidden>8192, non-2D). This removes the be kernel, its launchers (launch_be/launch_norm_be) and the OPUS_BE bucket table entirely.

The ctypes caller always appends the current stream to the call args, but the argtypes builder only declared that trailing hipStream_t when the op had a tensor param. For a torch-free ctypes module (all params non-tensor, e.g. the opus rmsnorm), argtypes ended up one short of the passed args, so ctypes took the variadic path (ffi_prep_cif_var). That is tolerated by some libffi builds but fails on others (observed as 'ffi_prep_cif_var failed' running vLLM + DeepSeek on the CI). Declare the stream unconditionally so the call is always non-variadic.

rmsnorm2d_fwd_with_add_dynamicquant / _with_add_smoothquant staged the residual with a torch residual_out.copy_(residual_in) before an in-place-add kernel. Give the quant kernel a residual_out pointer and a compile-time OOP flag (mirroring the norm kernel): OOP=true reads residual_in and writes residual_out in one pass; the OOP=false (in-place / no-add) instantiation is byte-identical to before, so those paths do not move. Removes the extra HBM copy -- 1.34x on add-dynamicquant at 7168. residual_in is now preserved (true out-of-place).

rmsnorm2d_fwd_with_dynamicquant(+add) still went to rmsnorm_quant_opus, which is ~2x slower than the ported arq kernel on the common shapes (e.g. 22.8us vs 11.6us at 4096x4096 int8). Route the common fp16/bf16, hidden<=8192, non-T5, 2-D case to the opus arq kernel (rmsnorm_quant / add_rmsnorm_quant, the bit-exact port of the HIP add_rmsnorm_quant_kernel) via a shared _use_arq_common helper; rmsnorm_quant_opus now handles only the fp32 / T5 / hidden>8192 fallback. Now at parity with HIP.

carlushuang added 30 commits July 2, 2026 08:44

Condense comments in opus rmsnorm sources

d9d184e

be-kernel: revert explicit fmaf (no fp16 gain); keep plain squared-sum

eda89d9

fp16's residual ~1-ulp vs CK is TU-context codegen (identical source is bit-exact in isolation), not the intra-thread FMA form. bf16 stays 100% bit-identical; fp16 is <=2 ulp. Reverted to the clearer expression.

rmsnorm: keep the launch-helper header as rmsnorm.h

32eb192

Rename csrc/include/rmsnorm_opus.h back to the original rmsnorm.h so the opus impl reuses the existing header name instead of adding an _opus one.

Merge remote-tracking branch 'origin/main' into carhuang/rmsnorm-opus

57d520e

rmsnorm opus: trim comments

bc96b4e

carlushuang added 7 commits July 4, 2026 09:44

rmsnorm opus quant: packed v_pk_mul_f32 square-sum + normalize

aece5e6

Match the reference's packed-f32 compute for the square-sum and the inv/weight multiplies. Lifts grouped small-n (n<=2048) a few points; neutral on the memory-bound shapes. Correctness unchanged (exact).

rmsnorm opus: refresh comments (opus is the sole quant backend)

f79051d

carlushuang requested a review from a team July 4, 2026 13:24

carlushuang added 10 commits July 4, 2026 14:13

carlushuang added 5 commits July 5, 2026 12:50

rmsnorm opus: condense comments across the split/kernel/py files

2277b3c

carlushuang mentioned this pull request Jul 5, 2026

[OPUS] RMSNorm backend using opus to reduce compile time (and keep feature / performance) #4059

Open

zufayu requested a review from junhaha666 July 6, 2026 02:01

carlushuang added 3 commits July 6, 2026 02:10

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[OPUS] Absorb module_rmsnorm_quant into the opus rmsnorm module#4080

[OPUS] Absorb module_rmsnorm_quant into the opus rmsnorm module#4080
carlushuang wants to merge 58 commits into
ROCm:mainfrom
carlushuang:carhuang/rmsnorm-opus-quant

carlushuang commented Jul 4, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

carlushuang commented Jul 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What is covered

Implementation notes

gfx1250

Feature coverage (all pass, gfx950 + gfx942)

Compile time (clean JIT build, gfx950)

Performance vs module_rmsnorm_quant — full sweep, both arches

Non-quant path (the hot path): what the engines actually call

Capability win: grouped quant at n = 5120 / 6144

n = model hidden_dim: which shapes matter

gfx942 add-path fix: bf16 store rounding

fused_qk_rmsnorm: no regression + strided rmsnorm fix

Why the small-n large-m add gap persists (root cause)

Impacted API

Uh oh!

github-actions Bot commented Jul 4, 2026

🏷️ CI Guide

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

carlushuang commented Jul 4, 2026 •

edited

Loading

Performance vs `module_rmsnorm_quant` — full sweep, both arches

`n` = model hidden_dim: which shapes matter