[OPUS] RMSNorm backend using opus to reduce compile time (and keep feature / performance)#4059
Merged
Merged
Conversation
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.
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
…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.
carlushuang
changed the title
- 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.
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.
…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' — this broke ATOM gpt-oss / DeepSeek / Kimi serving, which torch.compile the model (input_layernorm -> rmsnorm2d_fwd). 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 via @compile_ops. Eager unchanged.
…fx942 regression) The residual_out/no-quant bf16 stores used round-to-nearest-even, which has no hardware bf16 cvt on gfx942 and lowers to a ~6-op/element software sequence, making bf16 output/ residual stores ~2x slower there. The CK/HIP reference truncates; switching to truncate matches it exactly. gfx950 (hardware bf16 cvt) is unaffected. All op_tests 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).
… 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.
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; its be+generic kernels are the expensive ones to instantiate. Split launch_norm by input dtype into separate TUs (opus_norm_bf16/fp16/ fp32, each its own .cu; entrypoints dispatch the dtype code). Kernels/launcher unchanged -> identical kernels, no functional/perf change. Parallel compile ~1.65s (was ~2.6s); op_tests pass, torch.compile works, TUs compile for gfx1250.
…old coverage into test_rmsnorm2d test_rmsnorm2d.py / test_rmsnorm2dFusedAddQuant.py already exercise the opus backend (the public APIs route to opus). Remove the extra opus-specific test + compile bench; fold fp32 input and gemma_norm into test_rmsnorm2d so nothing is lost. All checks pass.
…for all sizes The be kernel was built to be bit-exact vs ck_tile, but that was never the common-case reference: main's rmsnorm2d_fwd dispatched bf16/fp16 non-T5 hidden<=8192 (the bucketed sizes) to add_rmsnorm_quant_kernel (HIP), and only fell back to ck_tile for T5 or hidden>8192. Verified bitwise: our generic kernel is <=1 ULP vs add_rmsnorm_quant_kernel -- identical closeness to what be gave -- so be provided no parity benefit, and perf is 96-100% of that kernel. be was purely the compile bottleneck. Remove rmsnorm_be_opus + launch_be + launch_norm_be + the OPUS_BE tile table; launch_norm now always uses the generic kernel. test_rmsnorm2d 810/810; parallel compile ~1.05s (was ~1.64s).
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). (cherry picked from commit c2900aa)
main dispatched the bf16/fp16, hidden<=8192, non-T5 case of rmsnorm2d_fwd / _with_add / _with_dynamicquant(+add) to the hand-tuned HIP add_rmsnorm_quant_kernel (module_rmsnorm_quant), and only used CK for the hidden>8192 / T5 fallback. This PR had routed everything through the opus kernels, which are 20-90% slower than the HIP kernel on those hot shapes -- a regression vs main. Restore main's dispatch: the common 2-D bf16/fp16 hidden<=8192 non-T5 path stays on the HIP module_rmsnorm_quant; opus now covers only what CK covered (fp32, T5, hidden>8192, non-2-D, gemma at any size). No perf regression, and the ~225s CK build is still gone.
vLLM calls aiter.rms_norm for the no-residual norm. It delegated rms_norm -> rmsnorm2d_fwd, so the hot path paid two torch_compile_guard layers plus the routing check -- ~1.47x the host dispatch cost of main's single-@compile_ops CK rms_norm at decode (m=1, n=7168: 11.7us vs 8.0us; the kernel itself, HIP arq, is at parity with CK ~7.5us). Route rms_norm and rmsnorm2d_fwd through one shared non-guarded helper so each pays a single guard (11.7us -> 10.4us). The remainder is guard + ctypes overhead that CUDA-graph capture (vLLM decode) elides entirely.
junhaha666
approved these changes
Jul 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
module_rmsnormhad a ~225s cold JIT build on gfx950 (#4055) — its CK backend emits 1360 TUs / 4972 template instantiations, most for quant/T5 variants the common path never uses.This replaces
module_rmsnormin place with a self-contained opus implementation (torch-free, ctypes, no pybind), cutting the cold build to ~1s while matching or beating the old kernels at runtime.Scope
module_rmsnorm— replaced. Deleted its CK (rmsnorm_ck_kernels.cu), the old vLLM HIP (rmsnorm_kernels.cu), pybind, and theRMSNORM_PYBINDmacro; the module now builds only the split opus translation units undercsrc/kernels/rmsnorm/. Module name unchanged.module_rmsnorm_quant— not touched. Separate, already-fast HIP module shared with other ops (all-reduce, MXFP pipelines); it owns grouped/MXFP4 quant (group_size/shuffle_scale, n ≤ 8192). The rmsnorm dispatch falls back to it only for those (and exotic dtypes), now behind an explicithidden <= 8192assert instead of a cryptic in-kernel abort. ([OPUS] Absorb module_rmsnorm_quant into the opus rmsnorm module #4080 absorbs it into opus.)Design
opus/hip_minimal.hpp+opus.hpp,-D__HIPCC_RTC__); no torch/ATen/pybind headers.rmsnorm_opus_kernel,rmsnorm_quant_opus); functional axes (fused-add / smooth / dynamic / save-unquant / T5) are runtime flags, so only in/out dtype × vector-width (and a compile-time gemma variant) are templated.rmsnorm_opus_norm_{bf16,fp16,fp32}.cu+rmsnorm_opus_norm_entry.cu+rmsnorm_opus_quant.cu) so the JIT compiles them in parallel.Coverage
Everything the old CK + HIP kernels did:
rms_norm,rms_norm_cu,fused_add_rms_norm_cu,rmsnorm2d_fwd(+_with_add), dynamic/smooth quant to int8/fp8 (+add, save-unquant), T5, andgemma_norm— for fp16/bf16/fp32 (fp32 is a net add), any hidden size (opus has no n>8192 cliff). Only grouped/MXFP4 quant (group_size/shuffle_scale), shared with other ops, stays inmodule_rmsnorm_quant(n ≤ 8192).Perf: before → after (gfx950 + gfx942)
Before = CK (
rmsnorm2d_fwd+ quant) + the vLLM HIP kernel (rms_norm/fused_add_rms_norm). After = opus. Higher = opus faster; HBM-bound at large hidden.rms_normruntimeImpacted aiter API
rms_norm,rmsnorm2d_fwd,rmsnorm2d_fwd_with_add,rms_norm_cu,fused_add_rms_norm_cu,rmsnorm2d_fwd_with_smoothquant(+add),rmsnorm2d_fwd_with_dynamicquant(+add)gemma_normargmodule_rmsnorm_quant, n ≤ 8192)rmsnorm2d_fwd_ck/_with_add_ck/_dynamicquant_ck/_with_add_dynamicquant_ck,gen_rms_norm_fake_tensormodule_rmsnorm_quant:rmsnorm_quant,add_rmsnorm_quant,add_rmsnorm,rmsnormshuffle_scale/ MXFP4)Testing
python op_tests/test_rmsnorm2d.py [--perf]andpython op_tests/test_rmsnorm2dFusedAddQuant.py— parity vs torch/CK across fp16/bf16/fp32,rmsnorm2d_fwd(+_with_add), T5 and gemma variants, dynamic/smooth quant, n up to 65536. (The standalonetest_rmsnorm_opus.py/bench_rmsnorm_compile.pywere folded into these two files.)