Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
carlushuang wants to merge 62 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
62 commits
Select commit Hold shift + click to select a range
54d0ebd
Add torch-free opus RMSNorm backend (module_rmsnorm_opus)
carlushuang Jul 2, 2026
d9d184e
Condense comments in opus rmsnorm sources
carlushuang Jul 2, 2026
c98be7c
Drop aiter_tensor.h: raw-pointer ctypes ABI + hip_minimal + __HIPCC_R…
carlushuang Jul 2, 2026
0f4bc9d
Cache row in registers to avoid double global read (single-pass)
carlushuang Jul 2, 2026
8430ca0
Full CK parity: opus dynamic/smooth quant (int8/fp8), T5, save-unquant
carlushuang Jul 2, 2026
0ec518e
Pack multiple rows per block for small hidden (2D block + segmented r…
carlushuang Jul 2, 2026
71200b8
Condense: merge plain + fused-add into one kernel, tighten comments
carlushuang Jul 2, 2026
6c85e66
Add opus mirrors of the CK entrypoints (symmetric _ck / _opus dispatch)
carlushuang Jul 2, 2026
b349bd2
Match CK exactly for fused-add: fp32 norm-input in default, rounded i…
carlushuang Jul 2, 2026
beddf58
Bit-exact vs CK for bf16 rmsnorm on the vn=8 tile buckets
carlushuang Jul 2, 2026
468d967
be-kernel: use scalar fp32 norm-input array (fp16 closer to CK)
carlushuang Jul 2, 2026
eda89d9
be-kernel: revert explicit fmaf (no fp16 gain); keep plain squared-sum
carlushuang Jul 2, 2026
fea17e5
Inline opus into rmsnorm.py (drop rmsnorm_opus.py); make opus the def…
carlushuang Jul 2, 2026
793a8d1
opus rmsnorm kernels: fp32 support + cleanups
carlushuang Jul 3, 2026
ac16131
rmsnorm: make opus the sole backend, remove CK
carlushuang Jul 3, 2026
47f3c44
rmsnorm: replace module_rmsnorm in place with the opus impl
carlushuang Jul 3, 2026
32eb192
rmsnorm: keep the launch-helper header as rmsnorm.h
carlushuang Jul 3, 2026
57d520e
Merge remote-tracking branch 'origin/main' into carhuang/rmsnorm-opus
carlushuang Jul 3, 2026
ccf704d
test_rmsnorm2d: pass use_model_sensitive_rmsnorm by keyword
carlushuang Jul 3, 2026
9111128
rmsnorm opus: implement gemma_norm, remove n>8192 quant fallback for it
carlushuang Jul 3, 2026
a20b9db
rmsnorm opus: carry gemma in fwd_traits instead of a separate kernel …
carlushuang Jul 4, 2026
34b3dd2
rmsnorm opus: drop the rmsnorm_opus namespace and rename kernels/traits
carlushuang Jul 4, 2026
2c2f3d8
rmsnorm opus: return launch geometry via std::pair + structured bindings
carlushuang Jul 4, 2026
fc1b117
rmsnorm opus: drop the 16-byte pointer-alignment gate
carlushuang Jul 4, 2026
9a47cef
rmsnorm opus: move host-only dtype vocabulary out of the kernel header
carlushuang Jul 4, 2026
2ddec16
rmsnorm opus: source element types from opus.hpp (host-includable sin…
carlushuang Jul 4, 2026
bc96b4e
rmsnorm opus: trim comments
carlushuang Jul 4, 2026
4ae9615
rmsnorm opus: derive qmax via aiter get_dtype_max instead of hardcoding
carlushuang Jul 4, 2026
5e81cec
rmsnorm opus: drop _use_opus dtype dispatch (opus is the only path)
carlushuang Jul 4, 2026
357ce0b
rmsnorm opus quant (phase 0): opus-clean detail header for grouped/MXFP4
carlushuang Jul 4, 2026
8ebee41
rmsnorm opus quant (phase 1): grouped/per-token quant kernel port
carlushuang Jul 4, 2026
85d427d
rmsnorm opus quant (phase 1): launcher + entrypoint + python wiring
carlushuang Jul 4, 2026
f4a6fcf
rmsnorm opus quant: coalesced vectorized IO (fix 10x perf regression)
carlushuang Jul 4, 2026
c635779
rmsnorm opus quant: DPP wave-reduce for the memory-light quant paths
carlushuang Jul 4, 2026
aece5e6
rmsnorm opus quant: packed v_pk_mul_f32 square-sum + normalize
carlushuang Jul 4, 2026
f79051d
rmsnorm opus: refresh comments (opus is the sole quant backend)
carlushuang Jul 4, 2026
fdce4dc
rmsnorm opus quant: cut block-reduce overhead on small-n quant paths
carlushuang Jul 4, 2026
35ae52d
rmsnorm opus quant: drop redundant leading barrier in block_reduce_1d
carlushuang Jul 4, 2026
de24c42
rmsnorm opus quant: fast reciprocals to erase the small-n divide over…
carlushuang Jul 4, 2026
3babdfa
rmsnorm opus quant: exact host 1/qmax + restore cu_num cache
carlushuang Jul 4, 2026
a824413
rmsnorm opus quant: real-DPP group-max reduce (warp_swap_)
carlushuang Jul 4, 2026
3ba8ca2
rmsnorm opus: accept non-contiguous input in rms_norm_opus
carlushuang Jul 4, 2026
8f8f541
rmsnorm opus quant: accept row-strided input in the quant entrypoints
carlushuang Jul 4, 2026
96e8649
rmsnorm opus quant: cheaper power-of-two arithmetic in the grouped sc…
carlushuang Jul 4, 2026
235bc8a
rmsnorm opus quant: fix grouped/fp4 crashes at non-2^k hidden dims
carlushuang Jul 4, 2026
9587f1e
rmsnorm opus quant: truncate fp32->bf16 stores to match the reference…
carlushuang Jul 4, 2026
65f9590
rmsnorm opus: thread an input row-stride through the 2d norm path
carlushuang Jul 4, 2026
25fdf01
rmsnorm opus: register public entrypoints as opaque custom ops for to…
carlushuang Jul 5, 2026
4f134fd
rmsnorm opus: out-of-place fused-add in one pass + be-kernel coverage…
carlushuang Jul 5, 2026
402e317
rmsnorm opus: split module_rmsnorm into per-feature compile units
carlushuang Jul 5, 2026
9115ff8
rmsnorm opus: split the norm TU by input dtype (bf16/fp16/fp32)
carlushuang Jul 5, 2026
2277b3c
rmsnorm opus: condense comments across the split/kernel/py files
carlushuang Jul 5, 2026
e8da6bb
rmsnorm opus: self-contained IO, stop touching aiter_opus_plus.h
carlushuang Jul 5, 2026
2e61a07
rmsnorm opus: drop test_rmsnorm_opus.py + bench_rmsnorm_compile.py, f…
carlushuang Jul 5, 2026
4279d28
rmsnorm opus: drop be kernel, route common case through opus arq
carlushuang Jul 5, 2026
accb6f2
jit: always declare the trailing hipStream_t in ctypes argtypes
carlushuang Jul 6, 2026
c2900aa
rmsnorm opus quant: out-of-place fused add, drop the host staging copy
carlushuang Jul 6, 2026
e13c6ca
rmsnorm opus quant: route per-token dynamicquant through the arq kernel
carlushuang Jul 6, 2026
fc3315c
Merge origin/main into carhuang/rmsnorm-opus-quant (4059 landed)
carlushuang Jul 6, 2026
74a0ceb
Merge remote-tracking branch 'origin/main' into carhuang/rmsnorm-opus…
carlushuang Jul 7, 2026
11b489c
rmsnorm opus: remove the now-dead module_rmsnorm_quant
carlushuang Jul 7, 2026
06c0421
Merge remote-tracking branch 'origin/main' into carhuang/rmsnorm-opus…
carlushuang Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion aiter/jit/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1078,7 +1078,6 @@ def _get_ck_exclude_modules():
"module_pa_metadata",
"module_ps_metadata",
"module_quant",
"module_rmsnorm_quant",
"module_rope_1c_uncached_fwd",
"module_rope_1c_uncached_bwd",
"module_rope_2c_uncached_fwd",
Expand Down
21 changes: 5 additions & 16 deletions aiter/jit/optCompilerConfig.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -686,7 +686,11 @@
"f'{AITER_CSRC_DIR}/kernels/rmsnorm/rmsnorm_opus_norm_fp16.cu'",
"f'{AITER_CSRC_DIR}/kernels/rmsnorm/rmsnorm_opus_norm_fp32.cu'",
"f'{AITER_CSRC_DIR}/kernels/rmsnorm/rmsnorm_opus_norm_entry.cu'",
"f'{AITER_CSRC_DIR}/kernels/rmsnorm/rmsnorm_opus_quant.cu'"
"f'{AITER_CSRC_DIR}/kernels/rmsnorm/rmsnorm_opus_quant.cu'",
"f'{AITER_CSRC_DIR}/kernels/rmsnorm/rmsnorm_opus_arq_i8.cu'",
"f'{AITER_CSRC_DIR}/kernels/rmsnorm/rmsnorm_opus_arq_fp8.cu'",
"f'{AITER_CSRC_DIR}/kernels/rmsnorm/rmsnorm_opus_arq_fp4.cu'",
"f'{AITER_CSRC_DIR}/kernels/rmsnorm/rmsnorm_opus_arq_entry.cu'"
],
"flags_extra_cc": [],
"flags_extra_hip": [
Expand All @@ -697,21 +701,6 @@
"verbose": "False",
"blob_gen_cmd": "''"
},
"module_rmsnorm_quant": {
"srcs": [
"f'{AITER_CSRC_DIR}/kernels/rmsnorm_quant_kernels.cu'",
"f'{AITER_CSRC_DIR}/pybind/rmsnorm_quant_pybind.cu'"
],
"flags_extra_cc": [],
"flags_extra_hip": [
"'-ffast-math'",
"f'-DOPUS_FP32_to_BF16_DEFAULT={os.environ.get(\"OPUS_FP32_to_BF16_DEFAULT\", 2)}'"
],
"extra_ldflags": "None",
"extra_include": [],
"verbose": "False",
"blob_gen_cmd": "''"
},
"module_fused_qk_rmsnorm_group_quant": {
"srcs": [
"f'{AITER_CSRC_DIR}/kernels/fused_qk_rmsnorm_group_quant.cu'",
Expand Down
Loading
Loading