diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index f84d7593141f2..6dcb758a234fe 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -113,6 +113,7 @@ Do not modify directly.* * com.microsoft.QuantizeLinear * com.microsoft.QuantizeWithOrder * com.microsoft.QuickGelu + * com.microsoft.RMSNormRotary * com.microsoft.Range * com.microsoft.ReduceSumInteger * com.microsoft.RelativePositionBias @@ -7098,6 +7099,71 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.RMSNormRotary** + + CUDA BF16 ONNX-arithmetic RMS normalization followed by non-interleaved partial RoPE. Input is [B,S,heads,stride]; a contiguous last-axis slice is normalized with FP32 scale. Normalization rounds to BF16 before rotation. Output is flattened [B,S,heads*width]. Width 256 and more than one row use one kernel for FP32 decode rotary or prefill; other shapes and BF16 decode rotary reuse existing CUDA kernels. Prefill mode uses reciprocal-multiply normalization and BF16 rotary intermediates; decode uses division normalization and the selected rotary arithmetic. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
epsilon : float
+
Finite nonnegative RMS epsilon.
+
flatten_batch : int
+
One flattens B*S to sequence and outputs [1,B*S,heads*width].
+
input_start : int
+
Nonnegative start of normalized last-axis slice.
+
input_width : int
+
Normalized width; zero uses the full last dimension.
+
rotary_embedding_dim : int (required)
+
Positive even rotary width, no larger than normalized width.
+
use_fp32_arithmetic : int
+
Decode rotary FP32 (1) or BF16 intermediates (0).
+
+ +#### Inputs (5 - 6) + +
+
X : T
+
BF16 input [B,S,heads,stride].
+
scale : F
+
FP32 RMS scale [width].
+
position_ids : I
+
One base position or [B,S] positions.
+
cosine : T
+
BF16 cache [cache_rows,rotary_embedding_dim/2].
+
sine : T
+
Same shape/type as cosine.
+
prefill_numerics (optional) : TI
+
Optional CPU int32 scalar 0/1.
+
+ +#### Outputs (1 - 2) + +
+
Y : T
+
Normalized and rotated [B,S,heads*width].
+
normalized_fp32 (optional) : F
+
Optional BF16-rounded normalized values widened to FP32, before rotation, [B,S,heads,width].
+
+ +#### Type Constraints + +
+
T : tensor(bfloat16)
+
BF16 data.
+
F : tensor(float)
+
FP32 scale and optional normalized output.
+
I : tensor(int64)
+
Position indices.
+
TI : tensor(int32)
+
Numerical control.
+
+ + ### **com.microsoft.Range** Creates a sequence of numbers that begins at `start` and extends by increments of `delta` diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 2497a7e6fd31a..6d084f06a97dc 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1164,6 +1164,7 @@ The **OpSet Version** column uses the following notation: |QuantizeLinear|*in* x:**T1**
*in* y_scale:**T1**
*in* y_zero_point:**T2**
*out* y:**T2**|1+|**T1** = tensor(float16)
**T2** = tensor(int8), tensor(uint8)| |QuantizeWithOrder|*in* input:**F**
*in* scale_input:**S**
*out* output:**Q**|1+|**F** = tensor(float), tensor(float16)
**Q** = tensor(int8)
**S** = tensor(float)| |QuickGelu|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| +|RMSNormRotary|*in* X:**T**
*in* scale:**F**
*in* position_ids:**I**
*in* cosine:**T**
*in* sine:**T**
*in* prefill_numerics:**TI**
*out* Y:**T**
*out* normalized_fp32:**F**|1+|**F** = tensor(float)
**I** = tensor(int64)
**T** = tensor(bfloat16)
**TI** = tensor(int32)| |RelativePositionBias|*in* bias_table:**T**
*in* query_length:**U**
*in* key_length:**U**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |RemovePadding|*in* input:**T**
*in* sequence_token_count:**M**
*out* output:**T**
*out* token_offset:**M**
*out* cumulated_seq_len:**M**
*out* max_seq_len:**M**|1+|**T** = tensor(float), tensor(float16)| |RestorePadding|*in* input:**T**
*in* token_offset:**M**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| diff --git a/docs/contrib_ops/cuda/rms_norm_rotary.md b/docs/contrib_ops/cuda/rms_norm_rotary.md new file mode 100644 index 0000000000000..2501a1f27c647 --- /dev/null +++ b/docs/contrib_ops/cuda/rms_norm_rotary.md @@ -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. diff --git a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc new file mode 100644 index 0000000000000..810016b96529c --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc @@ -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 + +#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()) + .TypeConstraint("F", DataTypeImpl::GetTensorType()) + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) + .InputMemoryType(OrtMemTypeCPUInput, 5), + RMSNormRotary); + +RMSNormRotary::RMSNormRotary(const OpKernelInfo& info) + : CudaKernel(info), + epsilon_(info.GetAttrOrDefault("epsilon", 1e-6f)), + input_start_(info.GetAttrOrDefault("input_start", 0)), + input_width_(info.GetAttrOrDefault("input_width", 0)), + rotary_dim_(info.GetAttrOrDefault("rotary_embedding_dim", 0)), + fp32_rotary_(info.GetAttrOrDefault("use_fp32_arithmetic", 1) == 1), + flatten_batch_(info.GetAttrOrDefault("flatten_batch", 0) == 1) { + const auto arithmetic = info.GetAttrOrDefault("use_fp32_arithmetic", 1); + const auto flatten = info.GetAttrOrDefault("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(0); + const auto* scale = context->Input(1); + const auto* positions = context->Input(2); + const auto* cosine = context->Input(3); + const auto* sine = context->Input(4); + const auto* mode = context->Input(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(SafeInt(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()[0] == 0 || mode->Data()[0] == 1), + "RMSNormRotary prefill_numerics must contain one int32 value 0 or 1"); + dense = mode->Data()[0] == 1; + } + const TensorShape norm_shape({shape[0], shape[1], shape[2], width}); + const int groups = SafeInt(norm_shape.Size() / width); + const int heads = SafeInt(shape[2]), stride = SafeInt(shape[3]); + const int sequence = SafeInt(output_sequence); + const int dim = SafeInt(width), rotary = SafeInt(rotary_dim_), cache_rows = SafeInt(cache[0]); + auto* output = context->Output(0, TensorShape({output_batch, output_sequence, SafeInt(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(input->Data()) + input_start_; + return LaunchRMSNormRotary( + Stream(context), x, scale->Data(), positions->Data(), + reinterpret_cast(cosine->Data()), + reinterpret_cast(sine->Data()), + reinterpret_cast(output->MutableData()), + normalized ? normalized->MutableData() : nullptr, + groups, heads, sequence, stride, dim, rotary, cache_rows, scalar, epsilon_, dense, fp32_rotary_); +} + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu new file mode 100644 index 0000000000000..6014e47983be1 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu @@ -0,0 +1,82 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/rms_norm_rotary.h" +#include +#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(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(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(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(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(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<<>>( + reinterpret_cast(input), scale, positions, + reinterpret_cast(cosine), reinterpret_cast(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 diff --git a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h new file mode 100644 index 0000000000000..6b71093b2b3e8 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// 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 diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index e8f435f99b4d7..5d349ddf1bade 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -176,6 +176,7 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, GatedRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, GatedRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, GatedRMSNorm); +class CUDA_MS_OP_CLASS_NAME(1, RMSNormRotary); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, BranchwiseRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, BranchwiseRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, BranchwiseRMSNorm); @@ -543,6 +544,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index b18836cf0b764..56d57cec1d700 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -6130,5 +6130,90 @@ ONNX_MS_OPERATOR_SET_SCHEMA( propagateShapeFromInputToOutput(ctx, 0, 0); })); +ONNX_MS_OPERATOR_SET_SCHEMA( + RMSNormRotary, 1, + OpSchema() + .SetDoc("CUDA BF16 ONNX-arithmetic RMS normalization followed by non-interleaved partial RoPE. " + "Input is [B,S,heads,stride]; a contiguous last-axis slice is normalized with FP32 scale. " + "Normalization rounds to BF16 before rotation. Output is flattened [B,S,heads*width]. " + "Width 256 and more than one row use one kernel for FP32 decode rotary or prefill; " + "other shapes and BF16 decode rotary reuse existing CUDA kernels. " + "Prefill mode uses reciprocal-multiply normalization and BF16 rotary intermediates; " + "decode uses division normalization and the selected rotary arithmetic.") + .Attr("epsilon", "Finite nonnegative RMS epsilon.", AttributeProto::FLOAT, 1e-6f) + .Attr("input_start", "Nonnegative start of normalized last-axis slice.", AttributeProto::INT, static_cast(0)) + .Attr("input_width", "Normalized width; zero uses the full last dimension.", AttributeProto::INT, static_cast(0)) + .Attr("rotary_embedding_dim", "Positive even rotary width, no larger than normalized width.", AttributeProto::INT) + .Attr("use_fp32_arithmetic", "Decode rotary FP32 (1) or BF16 intermediates (0).", AttributeProto::INT, static_cast(1)) + .Attr("flatten_batch", "One flattens B*S to sequence and outputs [1,B*S,heads*width].", AttributeProto::INT, static_cast(0)) + .Input(0, "X", "BF16 input [B,S,heads,stride].", "T") + .Input(1, "scale", "FP32 RMS scale [width].", "F") + .Input(2, "position_ids", "One base position or [B,S] positions.", "I") + .Input(3, "cosine", "BF16 cache [cache_rows,rotary_embedding_dim/2].", "T") + .Input(4, "sine", "Same shape/type as cosine.", "T") + .Input(5, "prefill_numerics", "Optional CPU int32 scalar 0/1.", "TI", OpSchema::Optional) + .Output(0, "Y", "Normalized and rotated [B,S,heads*width].", "T") + .Output(1, "normalized_fp32", + "Optional BF16-rounded normalized values widened to FP32, before rotation, " + "[B,S,heads,width].", + "F", OpSchema::Optional) + .TypeConstraint("T", {"tensor(bfloat16)"}, "BF16 data.") + .TypeConstraint("F", {"tensor(float)"}, "FP32 scale and optional normalized output.") + .TypeConstraint("I", {"tensor(int64)"}, "Position indices.") + .TypeConstraint("TI", {"tensor(int32)"}, "Numerical control.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + if (ctx.getNumOutputs() > 1) { + updateOutputElemType(ctx, 1, TensorProto::FLOAT); + } + const int64_t width_attr = getAttribute(ctx, "input_width", 0); + const int64_t start = getAttribute(ctx, "input_start", 0); + const int64_t rotary = getAttribute(ctx, "rotary_embedding_dim", 0); + const int64_t mode = getAttribute(ctx, "use_fp32_arithmetic", 1); + const int64_t flatten = getAttribute(ctx, "flatten_batch", 0); + if (width_attr < 0 || start < 0 || rotary <= 0 || rotary % 2 || (mode != 0 && mode != 1) || + (flatten != 0 && flatten != 1)) { + fail_shape_inference("RMSNormRotary has invalid slice, rotary width or mode."); + } + if (!hasInputShape(ctx, 0)) return; + const auto& input = getInputShape(ctx, 0); + if (input.dim_size() != 4) fail_shape_inference("RMSNormRotary input must have rank four."); + auto normalized = input; + if (width_attr) normalized.mutable_dim(3)->set_dim_value(width_attr); + if (input.dim(3).has_dim_value() && normalized.dim(3).has_dim_value() && + (normalized.dim(3).dim_value() <= 0 || start > input.dim(3).dim_value() || + normalized.dim(3).dim_value() > input.dim(3).dim_value() - start || + rotary > normalized.dim(3).dim_value())) { + fail_shape_inference("RMSNormRotary slice or rotary width exceeds input width."); + } + if (hasInputShape(ctx, 1)) { + const auto& scale = getInputShape(ctx, 1); + if (scale.dim_size() != 1 || (scale.dim(0).has_dim_value() && normalized.dim(3).has_dim_value() && + scale.dim(0).dim_value() != normalized.dim(3).dim_value())) { + fail_shape_inference("RMSNormRotary scale must have shape [width]."); + } + } + ONNX_NAMESPACE::TensorShapeProto output; + *output.add_dim() = input.dim(0); + *output.add_dim() = input.dim(1); + if (flatten) { + output.mutable_dim(0)->set_dim_value(1); + output.mutable_dim(1)->Clear(); + if (input.dim(0).has_dim_value() && input.dim(1).has_dim_value()) { + const auto batch = input.dim(0).dim_value(), sequence = input.dim(1).dim_value(); + if (batch > 0 && sequence > INT64_MAX / batch) fail_shape_inference("RMSNormRotary sequence overflow."); + output.mutable_dim(1)->set_dim_value(batch * sequence); + } + } + auto* flat = output.add_dim(); + if (input.dim(2).has_dim_value() && normalized.dim(3).has_dim_value()) { + const auto heads = input.dim(2).dim_value(), width = normalized.dim(3).dim_value(); + if (heads <= 0 || width > INT64_MAX / heads) fail_shape_inference("Invalid RMSNormRotary head width."); + flat->set_dim_value(heads * width); + } + updateOutputShape(ctx, 0, output); + if (ctx.getNumOutputs() > 1) updateOutputShape(ctx, 1, normalized); + })); + } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index bec1ef0fde030..d45e3d4f65b28 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -91,6 +91,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, PagedAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, LinearAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, LinearAttentionGate); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, BranchwiseRMSNorm); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RMSNormRotary); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ScaledSiLU); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPreMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPostMix); @@ -227,6 +228,7 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc b/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc new file mode 100644 index 0000000000000..73c32a8ac4878 --- /dev/null +++ b/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include + +#include "core/common/float8.h" +#include "gtest/gtest.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +namespace { + +template +std::vector ToTensorType(const std::vector& data) { + if constexpr (std::is_same_v) { + return ToFloat16(data); + } else if constexpr (std::is_same_v) { + return ToBFloat16(data); + } else { + return data; + } +} + +std::vector> CudaProvider() { + std::vector> providers; + if (auto cuda = DefaultCudaExecutionProvider()) { + providers.push_back(std::move(cuda)); + } + return providers; +} + +TEST(RMSNormRotaryTest, RMSNormRotaryExactUnfusedParity) { + for (const int64_t width : {64, 256, 320}) { + for (const int64_t sequence : {0, 1, 2, 7}) { + for (const bool packed : {false, true}) { + for (const int32_t dense : {0, 1}) { + for (const bool fp32 : {false, true}) { + const int64_t batch = packed || sequence == 7 ? 2 : 1; + const int64_t heads = sequence == 1 && !packed ? 1 : 2; + const int64_t stride = packed ? width * 2 : width; + const int64_t start = packed ? 3 : 0; + const int64_t rotary = packed ? 64 : width; + constexpr int64_t cache_rows = 64; + const int64_t output_batch = packed ? 1 : batch; + const int64_t output_sequence = packed ? batch * sequence : sequence; + SCOPED_TRACE("width=" + std::to_string(width) + " S=" + std::to_string(sequence) + + " packed=" + std::to_string(packed) + " dense=" + std::to_string(dense) + + " fp32=" + std::to_string(fp32)); + std::vector input(batch * sequence * heads * stride); + std::vector scale(width), normalized(batch * sequence * heads * width); + std::vector norm(normalized.size()), expected(normalized.size()); + std::vector cosine(cache_rows * rotary / 2), sine(cosine.size()); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = BFloat16(std::sin(static_cast(i) * .731f) * + (i % 19 == 0 ? 64.f : .03f * static_cast(i % 11))); + } + for (int64_t i = 0; i < width; ++i) scale[i] = .8f + .0013f * static_cast(i); + for (size_t i = 0; i < cosine.size(); ++i) { + cosine[i] = BFloat16(std::cos(static_cast(i) * .13f)); + sine[i] = BFloat16(std::sin(static_cast(i) * .13f)); + } + const std::vector input_shape{batch, sequence, heads, stride}; + const std::vector norm_shape{batch, sequence, heads, width}; + const std::vector output_shape{output_batch, output_sequence, heads * width}; + const bool scalar = !packed; + std::vector positions(scalar ? 1 : output_batch * output_sequence, + cache_rows - output_sequence); + if (!scalar) { + for (size_t i = 0; i < positions.size(); ++i) positions[i] = static_cast((i * 7 + 3) % cache_rows); + if (!positions.empty()) positions.back() = cache_rows - 1; + } + const auto position_shape = scalar ? std::vector{1} : std::vector{output_batch, output_sequence}; + const auto run = [](OpTester& test) { + auto providers = CudaProvider(); + ASSERT_FALSE(providers.empty()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + }; + if (CudaProvider().empty()) GTEST_SKIP(); + const int64_t groups = batch * sequence * heads; + for (int64_t g = 0; g < groups; ++g) { + const auto* x = input.data() + g * stride + start; + float partial[2] = {0.f, 0.f}; + for (int w = 0; w < 2; ++w) { + float warp_sum[32] = {}; + for (int lane = 0; lane < 32; ++lane) { + const int tid = w * 32 + lane; + float s = 0.f; + for (int col = tid; col < width; col += 64) { + const float v = x[col].ToFloat(); + s += v * v; + } + warp_sum[lane] = s; + } + for (int step = 16; step > 0; step >>= 1) { + for (int lane = 0; lane < step; ++lane) warp_sum[lane] += warp_sum[lane + step]; + } + partial[w] = warp_sum[0]; + } + const float total = partial[0] + partial[1]; + const float rms = std::sqrt((total / static_cast(width)) + 1e-6f); + const float inv_rms = 1.f / rms; + for (int64_t col = 0; col < width; ++col) { + const float v = x[col].ToFloat(); + const float n = dense ? (v * inv_rms) : (v / rms); + const BFloat16 rounded(n * scale[col]); + norm[g * width + col] = rounded; + normalized[g * width + col] = rounded.ToFloat(); + } + const int64_t token = g / heads; + const int64_t base = positions[scalar ? 0 : token]; + const bool valid = scalar ? (base >= 0 && base <= cache_rows - output_sequence) + : (base >= 0 && base < cache_rows); + const int64_t pos = valid && scalar ? base + token % output_sequence : base; + const int64_t half = rotary / 2; + for (int64_t col = 0; col < width; ++col) { + BFloat16 val = norm[g * width + col]; + if (col < rotary && valid) { + const int64_t partner = (col + half) % rotary; + const int64_t cache = pos * half + col % half; + const float a = val.ToFloat(); + const float b = norm[g * width + partner].ToFloat(); + const float c = cosine[cache].ToFloat(); + const float s = sine[cache].ToFloat(); + if (!dense && fp32) { + val = BFloat16((a * c) + ((col < half ? -b : b) * s)); + } else { + const float p1 = BFloat16(a * c).ToFloat(); + const float signed_b = BFloat16((col < half ? -1.f : 1.f) * b).ToFloat(); + const float p2 = BFloat16(signed_b * s).ToFloat(); + val = BFloat16(p1 + p2); + } + } + expected[g * width + col] = val; + } + } + OpTester fused("RMSNormRotary", 1, kMSDomain); + fused.AddAttribute("input_width", width); + fused.AddAttribute("input_start", start); + fused.AddAttribute("rotary_embedding_dim", rotary); + fused.AddAttribute("use_fp32_arithmetic", fp32 ? 1 : 0); + fused.AddAttribute("flatten_batch", packed ? 1 : 0); + fused.AddAttribute("epsilon", 1e-6f); + fused.AddInput("X", input_shape, input); + fused.AddInput("scale", {width}, scale); + fused.AddInput("position", position_shape, positions); + fused.AddInput("cos", {cache_rows, rotary / 2}, cosine); + fused.AddInput("sin", {cache_rows, rotary / 2}, sine); + fused.AddInput("mode", {1}, {dense}); + fused.AddOutput("Y", output_shape, expected); + fused.AddOutput("normalized", norm_shape, normalized); + fused.SetOutputTolerance(0.f, 0.f); + run(fused); + } + } + } + } + } +} + +TEST(RMSNormRotaryTest, RMSNormRotaryRejectsInvalidContracts) { + for (int mode = 0; mode < 8; ++mode) { + SCOPED_TRACE(mode); + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + OpTester test("RMSNormRotary", 1, kMSDomain); + test.AddAttribute("input_width", 256); + test.AddAttribute("input_start", mode == 0 ? 1 : 0); + test.AddAttribute("rotary_embedding_dim", mode == 1 ? 33 : 32); + if (mode == 2) test.AddAttribute("epsilon", -1.f); + if (mode == 3) test.AddAttribute("use_fp32_arithmetic", 2); + test.AddInput("X", {1, 2, 2, 256}, std::vector(1024)); + test.AddInput("scale", {mode == 4 ? 128 : 256}, std::vector(mode == 4 ? 128 : 256, 1.f)); + test.AddInput("pos", {mode == 5 ? 3 : 1}, std::vector(mode == 5 ? 3 : 1)); + test.AddInput("cos", {8, 16}, std::vector(128)); + test.AddInput("sin", {mode == 6 ? 7 : 8, 16}, std::vector(mode == 6 ? 112 : 128)); + test.AddInput("mode", {1}, {mode == 7 ? 2 : 0}); + test.AddOutput("Y", {1, 2, 512}, std::vector(1024)); + test.Run(OpTester::ExpectResult::kExpectFailure, + mode == 0 ? "slice" : mode == 1 ? "rotary" + : mode == 2 ? "requires finite" + : mode == 3 ? "invalid slice" + : mode == 4 ? "scale" + : mode == 5 ? "positions" + : mode == 6 ? "cosine/sine" + : "prefill_numerics", + {}, nullptr, &providers); + } +} + +} // namespace +} // namespace onnxruntime::test