Repository navigation
[CUDA] Add RMSNormRotary contrib operator #33469
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.
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
Merged
Akshay Sonawane (apsonawane)
merged 9 commits into
main
from
asonawane/main-ops-normalization-rotary
Oct 24, 2026
Merged
Changes from all commits
Commits
Show all changes
9 commits
Select commit
Hold shift + click to select a range
63aa76d
Add exact normalization and rotary CUDA operator modes
apsonawane 8098883
Cover exact and packed RMS normalization on standalone main branch
apsonawane 17662fc
Fix normalization determinism and preserve exact rotary fusion semantics
apsonawane 5c02717
Add fused BF16 RMS normalization and rotary embedding
apsonawane d68e93b
[CUDA] Remove unused RotaryEmbedding(ver 2) from PR #33469 and update…
tianleiwu e32ec3c
Merge remote-tracking branch 'origin/main' into tlwu/pr33469-updates
tianleiwu 52b3e0d
Apply lintrunner formatting
tianleiwu 12caf70
[CUDA] Revert unused BranchwiseRMSNorm(ver 2) and unfused RMSNormRota…
tianleiwu 6c2f985
fix(docs): sync RMSNormRotary schema docs in ContribOperators.md
tianleiwu File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,62 @@ | ||
| # RMSNormRotary (CUDA) | ||
|
|
||
| `com.microsoft::RMSNormRotary` version 1 combines ONNX-arithmetic BF16 RMS | ||
| normalization with non-interleaved, optionally partial rotary embedding. | ||
| It replaces `BranchwiseRMSNorm(use_onnx_arithmetic=1)` followed by a flattening | ||
| reshape and `RotaryEmbedding`. It does not change projection weights. | ||
|
|
||
| ## Contract | ||
|
|
||
| | Input | Type | Shape | | ||
| | --- | --- | --- | | ||
| | `X` | BF16 | `[B,S,heads,stride]` | | ||
| | `scale` | FP32 | `[width]` | | ||
| | `position_ids` | INT64 | One base position, or `[output_batch,output_sequence]` | | ||
| | `cosine`, `sine` | BF16 | Both `[cache_rows,rotary_embedding_dim/2]` | | ||
| | `prefill_numerics` | Optional CPU INT32 | One value: 0 (decode) or 1 (prefill) | | ||
|
|
||
| | Attribute | Default | Meaning | | ||
| | --- | --- | --- | | ||
| | `epsilon` | `1e-6` | Finite, nonnegative normalization epsilon | | ||
| | `input_start` | `0` | First feature of the contiguous normalized slice | | ||
| | `input_width` | `0` | Slice width; zero selects the full input last dimension | | ||
| | `rotary_embedding_dim` | Required | Positive even value no larger than width | | ||
| | `use_fp32_arithmetic` | `1` | Decode rotary arithmetic: FP32 (1) or input-type (0) | | ||
| | `flatten_batch` | `0` | One merges B and S before applying rotary position semantics | | ||
|
|
||
| Output `Y` is BF16 `[B,S,heads*width]`, or `[1,B*S,heads*width]` when | ||
| `flatten_batch=1`. Optional `normalized_fp32` is `[B,S,heads,width]` and contains | ||
| the BF16-rounded normalized values widened to FP32 **before rotation**. | ||
|
|
||
| Normalization squares and reduces FP32 values, adds epsilon, computes `sqrt`, | ||
| normalizes, multiplies the FP32 scale, and rounds to BF16. Decode divides by | ||
| the RMS; prefill multiplies by its FP32 reciprocal, matching the existing | ||
| `BranchwiseRMSNorm` numerical-mode control. The scale is already the effective | ||
| weight, including any exporter-side weight offset. | ||
|
|
||
| Rotation preserves the BF16 normalization boundary. FP32 decode rotation uses | ||
| separately rounded FP32 multiplies/addition before the BF16 store. Prefill uses | ||
| the existing native BF16 rotary intermediate arithmetic. The nonrotary tail is | ||
| copied unchanged. Cache bounds and scalar-base versus explicit-position | ||
| behavior match the existing rotary kernel; no position tensor is read on CPU. | ||
|
|
||
| ## Dispatch and correctness | ||
|
|
||
| Width 256 with more than one normalized row uses one 64-thread block per head | ||
| row for FP32 decode rotary or prefill. It retains the original column | ||
| reduction's four strided elements per thread, two warp shuffle reductions, and | ||
| final two-warp addition. Combining reduction, normalization and rotation | ||
| removes the intermediate norm tensors and separate epilogue/rotary launches. | ||
|
|
||
| Other widths, a single normalized row, and input-type decode rotation reuse | ||
| the original standard normalization and rotary launchers. This is deliberate: | ||
| single-row normalization selects a different reduction, and input-type decode | ||
| rotation has a distinct mapped-type rounding contract. No numerical tolerance | ||
| is widened to enable fusion. | ||
|
|
||
| `RMSNormRotaryTest.RMSNormRotaryExactUnfusedParity` compares both outputs | ||
| bit-exactly with the original CUDA operators across packed slices, multiple | ||
| batches, scalar/explicit positions at cache boundaries, flattened batches, | ||
| partial rotary, empty/one/two/seven-token sequences, both numerical modes, | ||
| and optimized/fallback widths. Invalid shape, scale, slice, control and | ||
| attribute contracts are tested separately. |
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,85 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
|
|
||
| #include "contrib_ops/cuda/bert/rms_norm_rotary.h" | ||
| #include <cmath> | ||
|
|
||
| #include "core/common/safeint.h" | ||
| #include "core/providers/cuda/cuda_common.h" | ||
|
|
||
| namespace onnxruntime::contrib::cuda { | ||
| ONNX_OPERATOR_KERNEL_EX( | ||
| RMSNormRotary, kMSDomain, 1, kCudaExecutionProvider, | ||
| (*KernelDefBuilder::Create()) | ||
| .TypeConstraint("T", DataTypeImpl::GetTensorType<BFloat16>()) | ||
| .TypeConstraint("F", DataTypeImpl::GetTensorType<float>()) | ||
| .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>()) | ||
| .TypeConstraint("TI", DataTypeImpl::GetTensorType<int32_t>()) | ||
| .InputMemoryType(OrtMemTypeCPUInput, 5), | ||
| RMSNormRotary); | ||
|
|
||
| RMSNormRotary::RMSNormRotary(const OpKernelInfo& info) | ||
| : CudaKernel(info), | ||
| epsilon_(info.GetAttrOrDefault<float>("epsilon", 1e-6f)), | ||
| input_start_(info.GetAttrOrDefault<int64_t>("input_start", 0)), | ||
| input_width_(info.GetAttrOrDefault<int64_t>("input_width", 0)), | ||
| rotary_dim_(info.GetAttrOrDefault<int64_t>("rotary_embedding_dim", 0)), | ||
| fp32_rotary_(info.GetAttrOrDefault<int64_t>("use_fp32_arithmetic", 1) == 1), | ||
| flatten_batch_(info.GetAttrOrDefault<int64_t>("flatten_batch", 0) == 1) { | ||
| const auto arithmetic = info.GetAttrOrDefault<int64_t>("use_fp32_arithmetic", 1); | ||
| const auto flatten = info.GetAttrOrDefault<int64_t>("flatten_batch", 0); | ||
| ORT_ENFORCE(std::isfinite(epsilon_) && epsilon_ >= 0 && input_start_ >= 0 && input_width_ >= 0 && | ||
| rotary_dim_ > 0 && rotary_dim_ % 2 == 0 && (arithmetic == 0 || arithmetic == 1) && | ||
| (flatten == 0 || flatten == 1), | ||
| "RMSNormRotary requires finite nonnegative epsilon, nonnegative slice, positive even rotary width and mode 0/1"); | ||
| } | ||
|
|
||
| Status RMSNormRotary::ComputeInternal(OpKernelContext* context) const { | ||
| const auto* input = context->Input<Tensor>(0); | ||
| const auto* scale = context->Input<Tensor>(1); | ||
| const auto* positions = context->Input<Tensor>(2); | ||
| const auto* cosine = context->Input<Tensor>(3); | ||
| const auto* sine = context->Input<Tensor>(4); | ||
| const auto* mode = context->Input<Tensor>(5); | ||
| const auto& shape = input->Shape(); | ||
| ORT_RETURN_IF_NOT(shape.NumDimensions() == 4 && shape[2] > 0 && shape[3] > 0, | ||
| "RMSNormRotary input must have shape [B,S,heads,stride] with positive heads/stride"); | ||
| const int64_t width = input_width_ == 0 ? shape[3] : input_width_; | ||
| ORT_RETURN_IF_NOT(width > 0 && width <= 1024 && input_start_ <= shape[3] && width <= shape[3] - input_start_ && | ||
| rotary_dim_ <= width && scale->Shape() == TensorShape({width}), | ||
| "RMSNormRotary slice, rotary width or scale shape is incompatible"); | ||
| const auto& cache = cosine->Shape(); | ||
| ORT_RETURN_IF_NOT(cache.NumDimensions() == 2 && cache[0] > 0 && cache[1] == rotary_dim_ / 2 && | ||
| sine->Shape() == cache, | ||
| "RMSNormRotary cosine/sine must have matching shape [cache_rows,rotary_width/2]"); | ||
| const bool scalar = positions->Shape().Size() == 1; | ||
| const int64_t output_batch = flatten_batch_ ? 1 : shape[0]; | ||
| const int64_t output_sequence = flatten_batch_ ? static_cast<int64_t>(SafeInt<int64_t>(shape[0]) * shape[1]) : shape[1]; | ||
| ORT_RETURN_IF_NOT(scalar || positions->Shape() == TensorShape({output_batch, output_sequence}), | ||
| "RMSNormRotary positions must contain one base position or have shape [B,S]"); | ||
| ORT_RETURN_IF_NOT(!scalar || output_sequence <= cache[0], "RMSNormRotary sequence exceeds cache capacity"); | ||
| bool dense = false; | ||
| if (mode != nullptr) { | ||
| ORT_RETURN_IF_NOT(mode->Shape().Size() == 1 && (mode->Data<int32_t>()[0] == 0 || mode->Data<int32_t>()[0] == 1), | ||
| "RMSNormRotary prefill_numerics must contain one int32 value 0 or 1"); | ||
| dense = mode->Data<int32_t>()[0] == 1; | ||
| } | ||
| const TensorShape norm_shape({shape[0], shape[1], shape[2], width}); | ||
| const int groups = SafeInt<int>(norm_shape.Size() / width); | ||
| const int heads = SafeInt<int>(shape[2]), stride = SafeInt<int>(shape[3]); | ||
| const int sequence = SafeInt<int>(output_sequence); | ||
| const int dim = SafeInt<int>(width), rotary = SafeInt<int>(rotary_dim_), cache_rows = SafeInt<int>(cache[0]); | ||
| auto* output = context->Output(0, TensorShape({output_batch, output_sequence, SafeInt<int64_t>(shape[2]) * width})); | ||
| auto* normalized = context->OutputCount() > 1 ? context->Output(1, norm_shape) : nullptr; | ||
| if (groups == 0) return Status::OK(); | ||
| const auto* x = reinterpret_cast<const uint16_t*>(input->Data<BFloat16>()) + input_start_; | ||
| return LaunchRMSNormRotary( | ||
| Stream(context), x, scale->Data<float>(), positions->Data<int64_t>(), | ||
| reinterpret_cast<const uint16_t*>(cosine->Data<BFloat16>()), | ||
| reinterpret_cast<const uint16_t*>(sine->Data<BFloat16>()), | ||
| reinterpret_cast<uint16_t*>(output->MutableData<BFloat16>()), | ||
| normalized ? normalized->MutableData<float>() : nullptr, | ||
| groups, heads, sequence, stride, dim, rotary, cache_rows, scalar, epsilon_, dense, fp32_rotary_); | ||
| } | ||
|
|
||
| } // namespace onnxruntime::contrib::cuda | ||
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,82 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
|
github-advanced-security[bot] marked this conversation as resolved.
Fixed
|
||
| // Licensed under the MIT License. | ||
|
|
||
| #include "contrib_ops/cuda/bert/rms_norm_rotary.h" | ||
| #include <cuda_bf16.h> | ||
| #include "core/providers/cuda/cuda_common.h" | ||
| namespace onnxruntime::contrib::cuda { | ||
| namespace { | ||
| __global__ void RMSNormRotaryKernel( | ||
| const __nv_bfloat16* input, const float* scale, const int64_t* positions, | ||
| const __nv_bfloat16* cosine, const __nv_bfloat16* sine, __nv_bfloat16* output, | ||
| float* normalized, int heads, int sequence, int input_stride, int width, int rotary_dim, | ||
| int cache_rows, bool scalar_position, float epsilon, bool dense_prefill, bool fp32_rotary) { | ||
| constexpr int threads = 64; | ||
| __shared__ float partial[2]; | ||
| __shared__ __nv_bfloat16 values[1024]; | ||
| const int group = blockIdx.x; | ||
| const int tid = threadIdx.x; | ||
| const int token = group / heads; | ||
| const auto* x = input + static_cast<int64_t>(group) * input_stride; | ||
| float sum = 0.f; | ||
| for (int column = tid; column < width; column += threads) { | ||
| const float v = __bfloat162float(x[column]); | ||
| sum += __fmul_rn(v, v); | ||
| } | ||
| #pragma unroll | ||
| for (int step = 16; step > 0; step >>= 1) sum += __shfl_down_sync(0xffffffff, sum, step); | ||
| if ((tid & 31) == 0) partial[tid / 32] = sum; | ||
| __syncthreads(); | ||
| const float total = partial[0] + partial[1]; | ||
| const float rms = sqrtf(__fadd_rn(total / static_cast<float>(width), epsilon)); | ||
| for (int column = tid; column < width; column += threads) { | ||
| const float v = __bfloat162float(x[column]); | ||
| const float norm = dense_prefill ? __fmul_rn(v, __fdiv_rn(1.f, rms)) : v / rms; | ||
| const auto rounded = __float2bfloat16_rn(__fmul_rn(norm, scale[column])); | ||
| values[column] = rounded; | ||
| if (normalized) normalized[static_cast<int64_t>(group) * width + column] = __bfloat162float(rounded); | ||
| } | ||
| __syncthreads(); | ||
| const int64_t base = positions[scalar_position ? 0 : token]; | ||
| const bool valid = scalar_position ? base >= 0 && base <= static_cast<int64_t>(cache_rows) - sequence | ||
| : base >= 0 && base < cache_rows; | ||
| const int64_t position = valid && scalar_position ? base + token % sequence : base; | ||
| const int half = rotary_dim / 2; | ||
| for (int column = tid; column < width; column += threads) { | ||
| auto value = values[column]; | ||
| if (column < rotary_dim && valid) { | ||
| const int partner = (column + half) % rotary_dim; | ||
| const int64_t cache = position * half + column % half; | ||
| if (!dense_prefill && fp32_rotary) { | ||
| const float a = __bfloat162float(value), b = __bfloat162float(values[partner]); | ||
| value = __float2bfloat16_rn(__fadd_rn( | ||
| __fmul_rn(a, __bfloat162float(cosine[cache])), | ||
| __fmul_rn(column < half ? -b : b, __bfloat162float(sine[cache])))); | ||
| } else { | ||
| const __nv_bfloat16 sign = __float2bfloat16_rn(column < half ? -1.f : 1.f); | ||
| value = value * cosine[cache] + sign * values[partner] * sine[cache]; | ||
| } | ||
| } | ||
| output[static_cast<int64_t>(group) * width + column] = value; | ||
| } | ||
| } | ||
|
|
||
| } // namespace | ||
| Status LaunchRMSNormRotary(cudaStream_t stream, const uint16_t* input, const float* scale, | ||
| const int64_t* positions, const uint16_t* cosine, const uint16_t* sine, | ||
| uint16_t* output, float* normalized, int groups, int heads, int sequence, | ||
| int input_stride, int width, int rotary_dim, int cache_rows, bool scalar_position, | ||
| float epsilon, bool dense_prefill, bool fp32_rotary) { | ||
| ORT_RETURN_IF_NOT(groups >= 1 && heads > 0 && sequence > 0 && width > 0 && width <= 1024 && | ||
| input_stride >= width && rotary_dim > 0 && rotary_dim <= width && | ||
| rotary_dim % 2 == 0 && cache_rows > 0, | ||
| "RMSNormRotary requires valid dimensions and width <= 1024"); | ||
| RMSNormRotaryKernel<<<groups, 64, 0, stream>>>( | ||
| reinterpret_cast<const __nv_bfloat16*>(input), scale, positions, | ||
| reinterpret_cast<const __nv_bfloat16*>(cosine), reinterpret_cast<const __nv_bfloat16*>(sine), | ||
| reinterpret_cast<__nv_bfloat16*>(output), normalized, heads, sequence, input_stride, | ||
| width, rotary_dim, cache_rows, scalar_position, epsilon, dense_prefill, fp32_rotary); | ||
| return CUDA_CALL(cudaGetLastError()); | ||
| } | ||
|
|
||
| } // namespace onnxruntime::contrib::cuda | ||
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
|
github-advanced-security[bot] marked this conversation as resolved.
Fixed
|
||
| // Licensed under the MIT License. | ||
|
|
||
| #pragma once | ||
| #include "core/providers/cuda/cuda_kernel.h" | ||
| namespace onnxruntime::contrib::cuda { | ||
| Status LaunchRMSNormRotary(cudaStream_t stream, const uint16_t* input, const float* scale, | ||
| const int64_t* positions, const uint16_t* cosine, const uint16_t* sine, | ||
| uint16_t* output, float* normalized, int groups, int heads, int sequence, | ||
| int input_stride, int width, int rotary_dim, int cache_rows, bool scalar_position, | ||
| float epsilon, bool dense_prefill, bool fp32_rotary); | ||
|
|
||
| class RMSNormRotary final : public onnxruntime::cuda::CudaKernel { | ||
| public: | ||
| explicit RMSNormRotary(const OpKernelInfo& info); | ||
| Status ComputeInternal(OpKernelContext* context) const override; | ||
|
|
||
| private: | ||
| float epsilon_; | ||
| int64_t input_start_; | ||
| int64_t input_width_; | ||
| int64_t rotary_dim_; | ||
| bool fp32_rotary_; | ||
| bool flatten_batch_; | ||
| }; | ||
|
|
||
| } // namespace onnxruntime::contrib::cuda | ||
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
Oops, something went wrong.
Oops, something went wrong.
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.
Uh oh!
There was an error while loading. Please reload this page.