feat: add fused rms_norm_rope operator (per-head RMSNorm + RoPE; CPU, NVIDIA) - #1567
Open
shsaihdsaiudh wants to merge 2 commits into
Open
shsaihdsaiudh wants to merge 2 commits into
shsaihdsaiudh wants to merge 2 commits into
Conversation
added 2 commits
August 24, 2026 19:39
Single-kernel fusion of the per-head q/k RMSNorm + RoPE chain used by attention blocks: 2 kernel launches + 2 memory round-trips per tensor become 1 in-place pass. - infiniop C layer: descriptor/workspace/calculate entry points with CPU + NVIDIA dispatch; info validation enforces full-rotary (head_dim == 2 x table_dim), pos I32/I64, sin/cos F32 contiguous - CUDA kernel: fp32 accumulation for the sum-of-squares, rsqrtf, rounds to data type first, then rotates -- replicating the standalone rope kernel's rounding path per variant (GPT_J/NEOX x half/bf16) - infinicore bridge: op::rms_norm_rope_ with graph-op support
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
1 participant
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.
Summary
Add a fused
rms_norm_ropeoperator: per-head RMSNorm + RoPE in a single kernel, consumed by InfiniLM's Qwen3 paged attention path (companion PR to InfiniLM, linked below).Note on layout: this branch is based on the pre-split monolithic layout (
src/infiniop/ops/...), developed before the unified component architecture refactor (#1406). The op itself (dispatcher, CPU kernel, NVIDIA CUDA kernel) is self-contained and ports mechanically to the new InfiniOps structure.Content
include/infiniop/ops/rms_norm_rope.h— C API (infiniopRMSNormRoPEDescriptor_t, create/plan/run)include/infinicore/ops/rms_norm_rope.hpp— C++ API (infinicore::op::rms_norm_rope_, in-place)src/infiniop/ops/rms_norm_rope/operator.cc— device dispatchsrc/infiniop/ops/rms_norm_rope/cpu/— CPU implementationsrc/infiniop/ops/rms_norm_rope/nvidia/+cuda/kernel.cuh— NVIDIA CUDA implementationSemantics: for each token,
x = rms_norm(x) * weightper head, then RoPE with the givenpos_idsand precomputed sin/cos tables. Full-rotary only (head_dim == 2 * table_dim).Platform status
Verification
dev_perf/op_check_rms_norm_rope.cppin the InfiniLM PR), results within fp tolerance.dev_perf/gap_analysis.mdv5/v6 in the InfiniLM PR.Reproduce
$INFINI_ROOT.python dev_perf/bench.py --engine infinilm --model Qwen/Qwen3-0.6B --attn-backend hybrid(fused) vs the unfused fallback path.dev_perf/op_check_rms_norm_rope.cppagainst the installed InfiniCore (commands in the file header).Companion InfiniLM PR: InfiniTensor/InfiniLM#585