add fused_qknorm hip kernel by XiaobingSuper · Pull Request #2442 · ROCm/aiter

XiaobingSuper · 2026-03-24T06:56:32Z

Motivation

Add a fused_qknorm hip kernel to improve kimi-2 model performance/

Technical Details

Fused two rmsnorm to a fused norm for a smaller number of tokens, which can improve decode performance/

Test Plan

python test_fused_qk_norm.py -d bf16 -m 2 4 8 16 32 64 128 256 512 1024 2048 4096 8192 16384 32768 65536 -n1 1536 -n2 512

Test Result

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Copilot

Pull request overview

This PR adds a new HIP fused Q/K RMSNorm kernel and wires it through the ROCm extension + Python API to improve decode performance for Kimi-2–like shapes (small token counts), along with a new benchmark script.

Changes:

Add fused_qk_rmsnorm HIP kernel (2D grid, Q and K computed in parallel).
Expose the kernel via pybind and a Python wrapper with a conditional fallback path.
Update JIT build config for the existing fused QK norm/RoPE/cache module and add a new op benchmark.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
op_tests/test_fused_qk_norm.py	Adds a benchmark/accuracy script for fused vs split RMSNorm on Q/K.
csrc/kernels/fused_qk_norm.cu	Implements the new fused QK RMSNorm HIP kernel + C++ entrypoint.
csrc/include/rocm_ops.hpp	Exposes `fused_qk_rmsnorm` via pybind in the existing fused QK norm/RoPE/cache module macro.
csrc/include/fused_qk_norm_rope_cache_quant.h	Declares the new `fused_qk_rmsnorm` API for pybind visibility.
aiter/ops/fused_qk_norm_rope_cache_quant.py	Adds the Python wrapper and `compile_ops` binding for `fused_qk_rmsnorm`.
aiter/jit/optCompilerConfig.json	Extends the module’s sources/includes/flags to compile and link the new kernel.

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

valarLip · 2026-03-24T12:02:46Z

nice, could you please help finish this refactor ROCm/ATOM#342

XiaobingSuper · 2026-03-24T12:13:12Z

nice, could you please help finish this refactor ROCm/ATOM#342

Yes, I will add a PR in ATOM side to add this fused_op based on [ROCm/ATOM#342].

valarLip · 2026-03-25T02:30:03Z

+    int q_in_stride  = q.stride(0);
+    int k_in_stride  = k.stride(0);
+
+    auto q_out = torch::empty({m, q_n}, q.options());


maybe expose output to user end, so we can have stride tensor for some opt cases

Added output option.

* add fused_qknorm hip kernel (#2442) * base impl * improve big M * optimize v2 * fallback split norm for bigger M * fix rebase failed * code format * fix include error * support outplace works for torch compile * add output args --------- Co-authored-by: Guanbao Yu <gyu@amd.com> * feat: add fast gelu (#2220) * add fast gelu * fix ut * refactor * add more log * change to empty to avoid perf downgrade * make black lint happy * make ruff happy * fix import error * Fix build CK pipeline (#2399) * Fix renamed CK pipeline * Fix CK version * Rename eight_warps to eight_waves for consistency --------- Co-authored-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> * [OPUS] gfx1250 support for opus wmma scale and moe_sorting kernel (#2449) * [OPUS] Add gfx1250 support for opus moe_sorting kernel Add gfx1250 (MI450) to the GPU architecture maps and fix the opus moe_sorting kernel to use the correct wait counter instruction on gfx12 ISA (s_wait_dscnt instead of s_waitcnt_lgkmcnt). * [OPUS] Add WMMA scale instruction support for gfx1250 Add scaled WMMA dispatch (BX32 int / BX16 long) to struct wmma for: - wmma_scale[16]_f32_16x16x128_f8f6f4 (fp8 fmt=0, fp4 fmt=4) - wmma_scale[16]_f32_32x16x128_f4 (dedicated fp4) Key changes in opus.hpp: - Add fmt_a/fmt_b format codes and per-lane E8M0 scale operator() overloads with compile-time scale_sel (OPSEL) for lane group selection - Fix vtype for packed fp4: use reg_bytes_a/b to compute correct hardware register size (elem * bits / 8 bytes, not elem * sizeof) - Add zero-pad helper for fp4-via-f8f6f4 (i32x8 -> i32x16) - Add scaled overloads to wmma_adaptor_swap_ab Device tests (op_tests/opus/device/test_wmma_scale.cu): - Raw warp-level: 6 variants (fp8/fp4 x BX32/BX16 + 32x16 fp4) - Tiled MMA: 1x1, 2x2, 4x1 wave configurations via make_tiled_mma - Per-lane random scale: E8M0 exponents in [122..133], bitwise exact - All tests PASS on gfx1250, SKIP on gfx942/gfx950 * [OPUS] Fix black/ruff lint: rename ambiguous var `l` to `lane` * reduce wasted get_module overhead for module with custom module name (#2455) * fix moe gemm tuned config (#2463) * fix * add e=256 k=8 tuned config * revert t=1024,2048 back to ck * update * update * Skip `test_metadata_redirect.py` on archs other than `gfx942` (#2456) This test was designed to run only on `gfx942`. We want to run all Triton tests on `gfx950` as well - this is the reason behind the proposed changes. * [Triton] Fix bench_mha (#2317) * fix bench_mha * fix bench_mha parsing logic of bench_attn_models * fix bench_mla_decode parsing logic of bench_attn_models --------- Co-authored-by: Bruno Mazzotti <bruno.mazzotti@amd.com> * test_mla_persistent.py split kv reference fix max_seq_q != 1 error (#2363) * CI: fix dubious ownership for sglang checkout (#2477) * CI: mark sglang checkout paths as safe directories Add safe.directory entries before dependency installation for both host and container checkout paths to avoid git dubious ownership failures on /sglang-checkout. * CI: align sglang safe.directory fix with upstream Move safe.directory configuration to run right after container startup and only mark /sglang-checkout, matching upstream sglang handling for dubious ownership. * CI: use pip editable install and safe.directory in runtime CI (#2474) * Fix CKTile blockscale GEMM to read strides from tensor metadata (#2466) The CKTile blockscale GEMM wrapper hardcoded leading-dimension strides as stride_A = K, stride_B = K, stride_C = N, assuming fully contiguous row-major layout. This produced silently wrong results when input tensors had non-standard strides. In vLLM on ROCm, _maybe_pad_fp8_weight pads FP8 weight tensors for alignment and then creates a narrowed view, producing tensors whose logical shape is [N, K] but whose physical stride is [K+pad, 1]. The hardcoded stride_B = K caused the kernel to read from wrong memory offsets, leading to garbage output. Fix: read leading-dimension strides from the PyTorch tensor metadata (XQ.stride(0), WQ.stride(0), etc.) instead of assuming dense layout. Add TORCH_CHECK assertions to verify inner-dimension contiguity (stride == 1), which is required by the CKTile kernel. The old CK backend (gemm_a8w8_blockscale_common.cuh) already reads strides from tensor metadata, which is why it was unaffected. Made-with: Cursor * Add LSE output support for MLA decode qseqlen=1 persistent kernel (gf… (#2440) * Add LSE output support for MLA decode qseqlen=1 persistent kernel (gfx950) * Add qseqlen fold for MLA on gfx950: use qh64 kernel instead of qh16 * Add qseqlen fold for MLA on gfx950: use qh64 kernel instead of qh16 --------- Co-authored-by: root <root@gbt350-odcdh2-a11-1.png-odc.dcgpu> * tuned qwen3.5 gemm (#2485) Signed-off-by: Guanbao Yu <gyu@amd.com> Co-authored-by: Guanbao Yu <gyu@amd.com> * [Triton] Flash Attention Triton Windows build support (#2433) * Initial FA-2 Triton Windows build support * Continue Work lint minimize diff windows smoke test address copilot fixes remove is_windows error improve windows message --------- Co-authored-by: 0xDELUXA <djernovevo@gmail.com> * fix(gemm): add EVEN_MN heuristic to restore vectorized store in gemm(#2482) * fix(gemm): add EVEN_MN heuristic to restore vectorized store in gemm_a16w16 The do_not_specialize=["M","N"] change (commit 3170a51) prevents kernel recompilation when M/N change, but removes tt.divisibility=16 from M/N, causing AxisInfo to lose contiguity through RemOp and CmpOp, which degrades buffer_store from vectorized dwordx2 (vec=4) to scalar short (vec=1). Introduce EVEN_MN constexpr heuristic that checks M%BLOCK_SIZE_M==0 and N%BLOCK_SIZE_N==0 at compile time. When true, skip the modulo wrap on offs_am/offs_bn and use unmasked tl.store, restoring contiguity for vectorized memory operations without sacrificing the recompilation benefit. Made-with: Cursor * fix(gemm): add EVEN_MN heuristic to batched_gemm_a8w8 kernel Same fix as gemm_a16w16: the do_not_specialize=["M","N"] removes tt.divisibility=16, breaking vectorized store/load. Add EVEN_MN constexpr heuristic to conditionally skip modulo and mask when M%BLOCK_SIZE_M==0 and N%BLOCK_SIZE_N==0, restoring contiguity. Made-with: Cursor --------- Co-authored-by: jianlian <jianlian@amd.com> * rm gemm_common bind (#2425) * rm gemm_commona= and quant type bind * retune failed shape in a8w8_bpreshuffle_tuned_gemm.csv * recover enum * assert when found duplicated tuned shape (#2376) * assert when found duplicated tuned shape * rm duplicated tuned shape and update tuned file in model_configs * fix lint --------- Signed-off-by: Guanbao Yu <gyu@amd.com> Co-authored-by: XiaobingZhang <xiaobingzhangupc@gmail.com> Co-authored-by: Guanbao Yu <gyu@amd.com> Co-authored-by: ChenYou <youchen@amd.com> Co-authored-by: Enrico Degregori <73224202+EnricoDeg@users.noreply.github.com> Co-authored-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> Co-authored-by: carlushuang <carlus.huang@amd.com> Co-authored-by: Elton <zhimding@amd.com> Co-authored-by: Bruno Mazzotti <bruno.mazzotti@amd.com> Co-authored-by: Michael Melesse <micmelesse@gmail.com> Co-authored-by: minmengdie <memin@amd.com> Co-authored-by: Xin Huang <Xin.Huang@amd.com> Co-authored-by: Sami Remes <samremes@amd.com> Co-authored-by: fangche123 <Fang.Che@amd.com> Co-authored-by: root <root@gbt350-odcdh2-a11-1.png-odc.dcgpu> Co-authored-by: Pleaplusone <ygan@amd.com> Co-authored-by: 0xDELUXA <djernovevo@gmail.com> Co-authored-by: jianhao <Jianhao.Liang@amd.com> Co-authored-by: jianlian <jianlian@amd.com> Co-authored-by: yzhou103 <Ying.Zhou2@amd.com>

XiaobingSuper requested review from a team and Copilot March 24, 2026 06:56

Copilot started reviewing on behalf of XiaobingSuper March 24, 2026 06:58

View session

Copilot AI reviewed Mar 24, 2026

View reviewed changes

XiaobingSuper force-pushed the xiaobing/fuse_qk_norm branch from ab05d72 to 93da489 Compare March 24, 2026 07:14

XiaobingSuper requested review from gbyu-amd and wuhuikx March 24, 2026 07:17

gbyu-amd requested review from ganyi1996ppo and xytpai March 24, 2026 07:23

wuhuikx requested a review from valarLip March 24, 2026 07:55

wuhuikx previously approved these changes Mar 25, 2026

View reviewed changes

valarLip reviewed Mar 25, 2026

View reviewed changes

XiaobingSuper and others added 9 commits March 24, 2026 22:09

base impl

75fba44

improve big M

d751f42

optimize v2

a689e0c

fallback split norm for bigger M

e4ab85b

fix rebase failed

e866e8f

code format

1c847e5

fix include error

463081c

support outplace works for torch compile

ed2dd95

add output args

19ac005

XiaobingSuper dismissed wuhuikx’s stale review via 19ac005 March 25, 2026 03:09

XiaobingSuper force-pushed the xiaobing/fuse_qk_norm branch from 91be957 to 19ac005 Compare March 25, 2026 03:09

XiaobingSuper requested a review from valarLip March 25, 2026 03:10

wuhuikx approved these changes Mar 25, 2026

View reviewed changes

valarLip approved these changes Mar 25, 2026

View reviewed changes

valarLip merged commit 149a8f6 into ROCm:main Mar 25, 2026
23 of 24 checks passed

rbrugaro-amd mentioned this pull request Apr 7, 2026

[ROCm] Add MLA dual RMS norm fusion (Q, KV) pass for DeepSeek/Kimi-K2 vllm-project/vllm#39242

Merged

5 tasks done

Navigation Menu

Repository navigation

add fused_qknorm hip kernel - #2442

add fused_qknorm hip kernel#2442
valarLip merged 9 commits into
ROCm:mainfrom
XiaobingSuper:xiaobing/fuse_qk_norm

XiaobingSuper commented Mar 24, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

valarLip commented Mar 24, 2026

Uh oh!

XiaobingSuper commented Mar 24, 2026

Uh oh!

valarLip Mar 25, 2026

Uh oh!

XiaobingSuper Mar 25, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

XiaobingSuper commented Mar 24, 2026

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

valarLip commented Mar 24, 2026

Uh oh!

XiaobingSuper commented Mar 24, 2026

Uh oh!

valarLip Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

XiaobingSuper Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants