From 63aa76df94330a6c82d9e173a068f8a5a690dff2 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 19:08:52 +0000 Subject: [PATCH 1/8] Add exact normalization and rotary CUDA operator modes --- docs/ContribOperators.md | 46 +++- .../cpu/bert/branchwise_rms_norm.cc | 23 +- .../cpu/bert/branchwise_rms_norm.h | 1 + .../cuda/bert/branchwise_rms_norm.cc | 65 +++++- .../cuda/bert/branchwise_rms_norm.h | 3 + .../cuda/bert/branchwise_rms_norm_impl.cu | 69 ++++++ .../cuda/bert/branchwise_rms_norm_impl.h | 7 + .../contrib_ops/cuda/bert/rotary_embedding.cc | 19 +- .../contrib_ops/cuda/bert/rotary_embedding.h | 1 + .../cuda/bert/rotary_embedding_impl.cu | 62 +++-- .../cuda/bert/rotary_embedding_impl.h | 6 +- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 12 + .../webgpu/bert/branchwise_rms_norm.cc | 6 +- .../core/graph/contrib_ops/bert_defs.cc | 213 ++++++++++++------ onnxruntime/core/graph/contrib_ops/ms_opset.h | 4 + .../cuda/reduction/reduction_functions.cu | 31 ++- .../cuda/reduction/reduction_functions.h | 5 + onnxruntime/core/session/environment.cc | 2 +- .../contrib_ops/rotary_embedding_op_test.cc | 145 ++++++++++++ .../core/graph/training_op_defs.cc | 2 +- 20 files changed, 593 insertions(+), 129 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 2703de7cc1fae..d1594969658f2 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -894,18 +894,39 @@ This version of the operator has been available since version 1 of the 'com.micr (..., C, H), or flattened shape (..., C * H) when num_branches is specified. The optional scale may have shape (C * H), (C, H), or (H). Arithmetic is performed in float32 and the result is converted to T. + With use_onnx_arithmetic=1, compute FP32 square, mean, epsilon addition, square + root, division, and scale multiplication with separate rounding boundaries. + CUDA reuses the standard fast ReduceMean summation order, including its scalar-result + path; it does not substitute reciprocal square root or change the reduction tree. + For multiple rows, FP32-rounded squaring is performed during the reduction load rather + than materializing a separate FP32 square tensor. + Optional Y_fp32 losslessly widens the already T-rounded Y and requires + use_onnx_arithmetic=1. WebGPU supports only the default mode without Y_fp32. + CUDA-only version 2 can select a contiguous last-axis interval directly from + packed rows using input_start and input_width. It retains the same reduction + order over the selected width and avoids materializing a Slice tensor. #### Version -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. +Version 1 retains the original input layout. Version 2 adds CUDA-only packed-input +selection; CPU and WebGPU implement version 1 only. #### Attributes
epsilon : float
Epsilon added before reciprocal square root.
-
num_branches : int
+
input_start : int (default is 0, version 2)
+
Nonnegative starting offset in the last input dimension. A nonzero offset requires input_width.
+
input_width : int (default is 0, version 2)
+
Width selected before normalization. Zero retains the full input. A positive width +requires use_onnx_arithmetic=1 and num_branches=0, and the selected interval must fit +within the last input dimension.
+
num_branches : int (default is 0)
Number of branches for flattened input. Omit or set to zero for grouped input.
+
use_onnx_arithmetic : int (default is 0)
+
Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. +Supported values are 0 (default) and 1.
#### Inputs (1 - 2) @@ -921,12 +942,17 @@ This version of the operator has been available since version 1 of the 'com.micr
Y : T
-
RMS-normalized output with the same shape as X.
+
RMS-normalized output with the same shape as X, except that version 2 with +positive input_width replaces the last dimension with input_width.
+
Y_fp32 (optional) : F
+
Float32 widening of the already T-rounded Y. Requires use_onnx_arithmetic=1.
#### Type Constraints
+
F : tensor(float)
+
Widened output is float32.
T : tensor(float), tensor(float16), tensor(bfloat16)
Constrain input and output to floating-point tensors.
M : tensor(float), tensor(float16), tensor(bfloat16)
@@ -7173,10 +7199,18 @@ This version of the operator has been available since version 1 of the 'com.micr RotaryEmbedding is the implementation of rotary positional embeddings (RoPE). The positions are represented as rotation matrices that are multiplied to query and key before the inner product of query and key is taken. + Version 2 adds optional FP32 arithmetic with activation-type storage. With + use_fp32_arithmetic=1, input and cache values are widened inside the CUDA kernel, + each multiplication and addition is rounded separately in float32, and only the + final result is converted back to the input type. This preserves a + Cast-to-float -> float RotaryEmbedding -> Cast-to-input-type chain. + Version 1 and the default arithmetic are unchanged. #### Version -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. +The latest version is 2 of the 'com.microsoft' operator set. Version 2 is currently +implemented by CUDA only; other providers retain version 1 and cannot silently +execute version 2 with legacy low-precision arithmetic. #### Attributes @@ -7191,6 +7225,10 @@ This version of the operator has been available since version 1 of the 'com.micr
Rotary embedding dimension. Default value is 0.
scale : float
Custom scale will be used if specified. Default value is 1.0
+
use_fp32_arithmetic : int (default is 0; version 2 only)
+
Widen input and cache values, perform separately rounded float32 multiplication +and addition, then round the output once to the input type. Supported values are +0 (legacy arithmetic) and 1.
#### Inputs diff --git a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc index 4b50b240bdea0..23e8b6b02b354 100644 --- a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc @@ -15,6 +15,7 @@ namespace onnxruntime::contrib { BranchwiseRMSNorm, kMSDomain, 1, T, kCpuExecutionProvider, \ KernelDefBuilder() \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("F", DataTypeImpl::GetTensorType()) \ .TypeConstraint("M", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}), \ @@ -47,7 +48,11 @@ template BranchwiseRMSNorm::BranchwiseRMSNorm(const OpKernelInfo& info) : LayerNormImpl(info, -1, info.GetAttrOrDefault("epsilon", 1e-5f), true), epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), - num_branches_(info.GetAttrOrDefault("num_branches", 0)) {} + num_branches_(info.GetAttrOrDefault("num_branches", 0)), + use_onnx_arithmetic_(info.GetAttrOrDefault("use_onnx_arithmetic", 0) != 0) { + const auto mode = info.GetAttrOrDefault("use_onnx_arithmetic", 0); + ORT_ENFORCE(mode == 0 || mode == 1, "use_onnx_arithmetic must be 0 or 1"); +} template Status BranchwiseRMSNorm::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, @@ -71,7 +76,10 @@ Status BranchwiseRMSNorm::Compute(OpKernelContext* context) const { } auto* y = context->Output(0, x->Shape()); - if (scale != nullptr && + auto* widened = context->OutputCount() > 1 ? context->Output(1, x->Shape()) : nullptr; + ORT_RETURN_IF_NOT(widened == nullptr || use_onnx_arithmetic_, + "Y_fp32 requires use_onnx_arithmetic=1"); + if (!use_onnx_arithmetic_ && scale != nullptr && scale->Shape().Size() == params.hidden && scale->GetElementType() == x->GetElementType()) { TensorShapeVector grouped_dims(x->Shape().GetDims().begin(), @@ -98,15 +106,20 @@ Status BranchwiseRMSNorm::Compute(OpKernelContext* context) const { const float value = static_cast(x_data[offset + h]); sum_sq += value * value; } - const float inv_rms = 1.0f / std::sqrt(sum_sq / static_cast(params.hidden) + epsilon_); + const float rms = std::sqrt(sum_sq / static_cast(params.hidden) + epsilon_); + const float inv_rms = 1.0f / rms; for (int64_t h = 0; h < params.hidden; ++h) { const int64_t scale_index = scale != nullptr && scale->Shape().Size() == params.hidden ? h : (group % params.branches) * params.hidden + h; const float weight = scale == nullptr ? 1.0f : LoadFloat(*scale, scale_index); - y_data[offset + h] = - static_cast(static_cast(x_data[offset + h]) * inv_rms * weight); + const float value = static_cast(x_data[offset + h]); + y_data[offset + h] = static_cast( + (use_onnx_arithmetic_ ? value / rms : value * inv_rms) * weight); + if (widened != nullptr) { + widened->MutableData()[offset + h] = static_cast(y_data[offset + h]); + } } }, 0); diff --git a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h index 63284794f3f87..bd58f86db058d 100644 --- a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h +++ b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h @@ -18,6 +18,7 @@ class BranchwiseRMSNorm final : public LayerNormImpl { private: float epsilon_; int64_t num_branches_; + bool use_onnx_arithmetic_; }; } // namespace onnxruntime::contrib diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc index ee21a71c2b4aa..1c91e9619364e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc @@ -9,24 +9,29 @@ #include "contrib_ops/cpu/hyper_connection_helper.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" +#include "core/providers/cuda/reduction/reduction_functions.h" namespace onnxruntime::contrib::cuda { using namespace onnxruntime::cuda; -#define REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(T) \ +#define REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(V, T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ - BranchwiseRMSNorm, kMSDomain, 1, T, kCudaExecutionProvider, \ + BranchwiseRMSNorm, kMSDomain, V, T, kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("F", DataTypeImpl::GetTensorType()) \ .TypeConstraint("M", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}), \ BranchwiseRMSNorm); -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(float) -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(MLFloat16) -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(BFloat16) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(1, float) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(1, MLFloat16) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(1, BFloat16) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(2, float) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(2, MLFloat16) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(2, BFloat16) #undef REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL @@ -34,30 +39,70 @@ template BranchwiseRMSNorm::BranchwiseRMSNorm(const OpKernelInfo& info) : CudaKernel(info), epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), - num_branches_(info.GetAttrOrDefault("num_branches", 0)) {} + num_branches_(info.GetAttrOrDefault("num_branches", 0)), + use_onnx_arithmetic_(info.GetAttrOrDefault("use_onnx_arithmetic", 0) != 0), + input_start_(info.GetAttrOrDefault("input_start", 0)), + input_width_(info.GetAttrOrDefault("input_width", 0)) { + const auto mode = info.GetAttrOrDefault("use_onnx_arithmetic", 0); + ORT_ENFORCE(mode == 0 || mode == 1, "use_onnx_arithmetic must be 0 or 1"); + ORT_ENFORCE(input_start_ >= 0 && input_width_ >= 0, "input_start and input_width must be nonnegative"); + ORT_ENFORCE(input_width_ != 0 || input_start_ == 0, "input_start requires input_width"); + ORT_ENFORCE(input_width_ == 0 || (use_onnx_arithmetic_ && num_branches_ == 0), + "input_width requires use_onnx_arithmetic=1 and grouped input"); +} template Status BranchwiseRMSNorm::ComputeInternal(OpKernelContext* context) const { using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; const auto* x = context->Input(0); const auto* scale = context->Input(1); + auto output_dims = x->Shape().AsShapeVector(); + ORT_RETURN_IF_NOT(!output_dims.empty(), "BranchwiseRMSNorm requires a non-scalar input"); + const int64_t input_stride = output_dims.back(); + if (input_width_ != 0) { + ORT_RETURN_IF_NOT(input_start_ <= input_stride && input_width_ <= input_stride - input_start_, + "input slice exceeds the last dimension"); + output_dims.back() = input_width_; + } + const TensorShape output_shape(output_dims); hyper_connection::StreamShape params; - ORT_RETURN_IF_ERROR(hyper_connection::ResolveStreamShape(x->Shape(), num_branches_, params)); + ORT_RETURN_IF_ERROR(hyper_connection::ResolveStreamShape(output_shape, num_branches_, params)); if (scale != nullptr) { ORT_RETURN_IF_ERROR(hyper_connection::ValidateScale(scale->Shape(), params)); } ORT_RETURN_IF_NOT(params.branches <= std::numeric_limits::max() && - params.hidden <= std::numeric_limits::max(), + params.hidden <= std::numeric_limits::max() && + (input_width_ == 0 || input_stride <= std::numeric_limits::max()), "branch or hidden dimension is too large for CUDA"); - const int64_t groups = x->Shape().Size() / params.hidden; - auto* y = context->Output(0, x->Shape()); + const int64_t groups = output_shape.Size() / params.hidden; + auto* y = context->Output(0, output_shape); + auto* widened = context->OutputCount() > 1 ? context->Output(1, output_shape) : nullptr; + ORT_RETURN_IF_NOT(widened == nullptr || use_onnx_arithmetic_, + "Y_fp32 requires use_onnx_arithmetic=1"); if (groups == 0) { return Status::OK(); } const auto* x_data = reinterpret_cast(x->Data()); auto* y_data = reinterpret_cast(y->MutableData()); + if (use_onnx_arithmetic_) { + ORT_RETURN_IF_NOT(groups <= std::numeric_limits::max(), + "ONNX arithmetic requires the row count to fit in int32"); + const auto buffer_size = compute_reduce_matrix_columns_buffer_size( + static_cast(groups), static_cast(params.hidden)); + auto squares = GetScratchBuffer(groups == 1 ? output_shape.Size() : 0, GetComputeStream(context)); + auto sums = GetScratchBuffer(groups, GetComputeStream(context)); + auto buffer = GetScratchBuffer(buffer_size, GetComputeStream(context)); + return LaunchStandardBranchwiseRMSNorm( + Stream(context), x_data + input_start_, scale == nullptr ? nullptr : scale->DataRaw(), + scale == nullptr ? x->GetElementType() : scale->GetElementType(), y_data, + static_cast(groups), static_cast(params.branches), static_cast(params.hidden), + scale == nullptr || scale->Shape().Size() == params.hidden, epsilon_, + widened == nullptr ? nullptr : widened->MutableData(), + squares.get(), sums.get(), buffer.get(), buffer_size, + static_cast(input_width_ == 0 ? params.hidden : input_stride)); + } return LaunchMixedScaleBranchwiseRMSNorm( Stream(context), x_data, scale == nullptr ? nullptr : scale->DataRaw(), scale == nullptr ? x->GetElementType() : scale->GetElementType(), y_data, diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h index 959fa12b4a46a..c26a8d47b41ff 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h @@ -16,6 +16,9 @@ class BranchwiseRMSNorm final : public onnxruntime::cuda::CudaKernel { private: float epsilon_; int64_t num_branches_; + bool use_onnx_arithmetic_; + int64_t input_start_; + int64_t input_width_; }; } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu index a56e55ce21331..240b057b827e0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu @@ -7,10 +7,12 @@ #include #include #include +#include #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" #include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/reduction/reduction_functions.h" namespace onnxruntime::contrib::cuda { namespace { @@ -29,6 +31,34 @@ __device__ __forceinline__ float LoadScale(const void* data, int64_t index, int return __bfloat162float(static_cast(data)[index]); } +template +__global__ void BranchwiseSquareKernel(const T* x, float* squares, int64_t count) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) { + const float value = to_float(x[index]); + squares[index] = __fmul_rn(value, value); + } +} + +template +__global__ void BranchwiseStandardEpilogueKernel( + const T* x, const void* scale, int scale_type, const float* sums, T* y, float* widened, + int64_t count, int branches, int hidden, bool shared_scale, float epsilon, int input_row_stride) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) { + const int64_t group = index / hidden; + const int64_t scale_index = shared_scale ? index % hidden : index % (static_cast(branches) * hidden); + const float weight = scale == nullptr ? 1.0f : LoadScale(scale, scale_index, scale_type); + const float mean = sums[group] / static_cast(hidden); + const float rms = sqrtf(__fadd_rn(mean, epsilon)); + const int64_t source_index = Packed ? group * input_row_stride + index % hidden : index; + const float normalized = to_float(x[source_index]) / rms; + const T rounded = from_float(__fmul_rn(normalized, weight)); + y[index] = rounded; + if (widened != nullptr) widened[index] = to_float(rounded); + } +} + template __global__ void MixedScaleBranchwiseRMSNormKernel(const T* x, const void* scale, int scale_type, T* y, int64_t groups, int branches, int hidden, @@ -77,6 +107,37 @@ __global__ void MixedScaleBranchwiseRMSNormKernel(const T* x, const void* scale, } // namespace +template +Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const void* scale, + int scale_type, T* y, int groups, int branches, + int hidden, bool shared_scale, float epsilon, + float* widened, float* squares, float* sums, void* reduction_buffer, + size_t reduction_buffer_size, int input_row_stride) { + const int64_t count = static_cast(groups) * hidden; + const int64_t blocks = (count + kThreads - 1) / kThreads; + ORT_RETURN_IF_NOT(blocks <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); + // Standard ReduceMean removes singleton dimensions and prefers row reduction for a scalar result. + if (groups == 1) { + BranchwiseSquareKernel<<(blocks), kThreads, 0, stream>>>(x, squares, count); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + ORT_RETURN_IF_ERROR((onnxruntime::cuda::reduce_matrix_rows( + stream, squares, sums, hidden, 1))); + } else { + using ReduceT = std::conditional_t, BFloat16, T>; + ORT_RETURN_IF_ERROR((onnxruntime::cuda::reduce_square_matrix_columns( + stream, reinterpret_cast(x), sums, groups, hidden, + reduction_buffer, reduction_buffer_size, input_row_stride))); + } + if (input_row_stride == hidden) { + BranchwiseStandardEpilogueKernel<<(blocks), kThreads, 0, stream>>>( + x, scale, scale_type, sums, y, widened, count, branches, hidden, shared_scale, epsilon, input_row_stride); + } else { + BranchwiseStandardEpilogueKernel<<(blocks), kThreads, 0, stream>>>( + x, scale, scale_type, sums, y, widened, count, branches, hidden, shared_scale, epsilon, input_row_stride); + } + return CUDA_CALL(cudaGetLastError()); +} + template Status LaunchMixedScaleBranchwiseRMSNorm(cudaStream_t stream, const T* x, const void* scale, int scale_type, T* y, int64_t groups, int branches, @@ -107,4 +168,12 @@ INSTANTIATE(__nv_bfloat16) #undef INSTANTIATE +#define INSTANTIATE_STANDARD(T) \ + template Status LaunchStandardBranchwiseRMSNorm(cudaStream_t, const T*, const void*, int, T*, \ + int, int, int, bool, float, float*, float*, float*, void*, size_t, int); +INSTANTIATE_STANDARD(float) +INSTANTIATE_STANDARD(half) +INSTANTIATE_STANDARD(__nv_bfloat16) +#undef INSTANTIATE_STANDARD + } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h index 549de28e62bff..a037130d825a2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h @@ -14,4 +14,11 @@ Status LaunchMixedScaleBranchwiseRMSNorm(cudaStream_t stream, const T* x, const int scale_type, T* y, int64_t groups, int branches, int hidden, bool shared_scale, float epsilon); +template +Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const void* scale, + int scale_type, T* y, int groups, int branches, + int hidden, bool shared_scale, float epsilon, + float* widened, float* squares, float* sums, void* reduction_buffer, + size_t reduction_buffer_size, int input_row_stride); + } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc index 0b38c6f0e5484..722e99b096418 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc @@ -15,11 +15,11 @@ namespace onnxruntime { namespace contrib { namespace cuda { -#define REGISTER_KERNEL_TYPED(T) \ +#define REGISTER_KERNEL_TYPED(T, Version) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ RotaryEmbedding, \ kMSDomain, \ - 1, \ + Version, \ T, \ kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ @@ -27,9 +27,12 @@ namespace cuda { .TypeConstraint("M", DataTypeImpl::GetTensorType()), \ RotaryEmbedding); -REGISTER_KERNEL_TYPED(float) -REGISTER_KERNEL_TYPED(MLFloat16) -REGISTER_KERNEL_TYPED(BFloat16) +REGISTER_KERNEL_TYPED(float, 1) +REGISTER_KERNEL_TYPED(MLFloat16, 1) +REGISTER_KERNEL_TYPED(BFloat16, 1) +REGISTER_KERNEL_TYPED(float, 2) +REGISTER_KERNEL_TYPED(MLFloat16, 2) +REGISTER_KERNEL_TYPED(BFloat16, 2) template RotaryEmbedding::RotaryEmbedding(const OpKernelInfo& info) : CudaKernel(info) { @@ -38,6 +41,9 @@ RotaryEmbedding::RotaryEmbedding(const OpKernelInfo& info) : CudaKernel(info) num_heads = static_cast(info.GetAttrOrDefault("num_heads", 0)); interleaved = (info.GetAttrOrDefault("interleaved", 0) == 1); is_packed_batching = (info.GetAttrOrDefault("is_packed_batching", 0) == 1); + const auto mode = info.GetAttrOrDefault("use_fp32_arithmetic", 0); + ORT_ENFORCE(mode == 0 || mode == 1, "use_fp32_arithmetic must be 0 or 1"); + use_fp32_arithmetic = mode == 1; } template @@ -83,7 +89,8 @@ Status RotaryEmbedding::ComputeInternal(OpKernelContext* context) const { parameters.position_ids_format, interleaved, device_prop.maxThreadsPerBlock, - parameters.transposed); + parameters.transposed, + use_fp32_arithmetic); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h index 55f654a2cea8f..4487746579a9a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h @@ -23,6 +23,7 @@ class RotaryEmbedding final : public CudaKernel { int rotary_embedding_dim; bool interleaved; bool is_packed_batching; + bool use_fp32_arithmetic; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu index 47c46b73f7b83..55648cc944249 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu @@ -10,6 +10,7 @@ Kernel implementation for rotary embeddings. #include "contrib_ops/cuda/bert/rotary_embedding_impl.h" #include +#include #include "core/common/safeint.h" #include "core/providers/cuda/cu_inc/common.cuh" @@ -20,7 +21,7 @@ namespace onnxruntime { namespace contrib { namespace cuda { -template +template __global__ void RotaryEmbeddingBSNH(T* output, // BxSxNxH const T* input, // BxSxNxH const T* cos_cache, // Mx(H/2) @@ -128,7 +129,14 @@ __global__ void RotaryEmbeddingBSNH(T* output, // BxSxNx } // Use values from shared memory - if constexpr (use_smem) { + if constexpr (use_fp32_arithmetic) { + const float a = static_cast(use_smem ? smem[i] : input_data[i]); + const float b = static_cast(use_smem ? smem[j] : input_data[j]); + const float cosine = static_cast(cos_data[cache_idx]); + const float sine = static_cast(sin_data[cache_idx]); + const float signed_b = sign < T(0) ? -b : b; + output_data[i] = T(__fadd_rn(__fmul_rn(a, cosine), __fmul_rn(signed_b, sine))); + } else if constexpr (use_smem) { output_data[i] = smem[i] * cos_data[cache_idx] + sign * smem[j] * sin_data[cache_idx]; } else { output_data[i] = input_data[i] * cos_data[cache_idx] + sign * input_data[j] * sin_data[cache_idx]; @@ -142,7 +150,8 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, - const int max_threads_per_block, const bool is_input_bnsh_format) { + const int max_threads_per_block, const bool is_input_bnsh_format, + const bool use_fp32_arithmetic) { int sequence_stride; int batch_stride; int heads_stride; @@ -168,7 +177,7 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu rotary_embedding_dim, max_sequence_length, position_ids_format, interleaved, max_threads_per_block, - in_strides, out_strides); + in_strides, out_strides, use_fp32_arithmetic); } template @@ -179,8 +188,8 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, const int max_threads_per_block, - int4 in_strides, int4 out_strides // strides in bnsh coord -) { + int4 in_strides, int4 out_strides, // strides in bnsh coord + const bool use_fp32_arithmetic) { // Note: Current implementation assumes head_size <= max_threads_per_block // because head_size is currently large for LLaMA-2. For smaller head_size // and num_heads values, we can create a block as `block(num_heads, head_size, 1)` @@ -196,20 +205,25 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu assert(head_size <= max_threads_per_block); - if (output == input) { - // In-place operation: use shared memory to avoid read-after-write hazards - size_t smem_size = head_size * sizeof(T); - RotaryEmbeddingBSNH<<>>( - output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, - num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, - interleaved, in_strides, out_strides); - } else { - // Separate buffers: no shared memory needed - RotaryEmbeddingBSNH<<>>( - output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, - num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, - interleaved, in_strides, out_strides); - } + auto launch = [&](auto mode) { + constexpr bool fp32 = decltype(mode)::value; + if (output == input) { + size_t smem_size = head_size * sizeof(T); + RotaryEmbeddingBSNH<<>>( + output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, + num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, + interleaved, in_strides, out_strides); + } else { + RotaryEmbeddingBSNH<<>>( + output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, + num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, + interleaved, in_strides, out_strides); + } + }; + if (use_fp32_arithmetic) + launch(std::true_type{}); + else + launch(std::false_type{}); return CUDA_CALL(cudaGetLastError()); } @@ -220,7 +234,8 @@ template Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, float* o const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, - const int max_threads_per_block, const bool is_input_bnsh_format); + const int max_threads_per_block, const bool is_input_bnsh_format, + const bool use_fp32_arithmetic); template Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, half* output, const half* input, const int64_t* position_ids, const int* past_sequence_lengths, const half* cos_cache, @@ -228,14 +243,15 @@ template Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, half* out const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, - const int max_threads_per_block, const bool is_input_bnsh_format); + const int max_threads_per_block, const bool is_input_bnsh_format, + const bool use_fp32_arithmetic); template Status LaunchRotaryEmbeddingKernel( cudaStream_t stream, BFloat16* output, const BFloat16* input, const int64_t* position_ids, const int* past_sequence_lengths, const BFloat16* cos_cache, const BFloat16* sin_cache, const int batch_size, const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, const int max_threads_per_block, - const bool is_input_bnsh_format); + const bool is_input_bnsh_format, const bool use_fp32_arithmetic); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h index 2d81e6b4067af..4422572fc271d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h @@ -27,7 +27,8 @@ Status LaunchRotaryEmbeddingKernel( const int position_ids_format, const bool interleaved, const int max_threads_per_block, - const bool is_input_bnsh_format); + const bool is_input_bnsh_format, + const bool use_fp32_arithmetic = false); template Status LaunchRotaryEmbeddingKernel( @@ -48,7 +49,8 @@ Status LaunchRotaryEmbeddingKernel( const bool interleaved, const int max_threads_per_block, int4 in_strides, - int4 out_strides); + int4 out_strides, + const bool use_fp32_arithmetic = false); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 546b5b12e178f..fe140646872de 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -160,6 +160,9 @@ class CUDA_ONNX_OP_TYPED_CLASS_NAME(1, MLFloat16, ParametricSoftplus); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, RotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, RotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, RotaryEmbedding); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, float, RotaryEmbedding); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, MLFloat16, RotaryEmbedding); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, BFloat16, RotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MRotaryEmbedding); @@ -179,6 +182,9 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, GatedRMSNorm); 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); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, float, BranchwiseRMSNorm); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, MLFloat16, BranchwiseRMSNorm); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, BFloat16, BranchwiseRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, ScaledSiLU); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, ScaledSiLU); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, ScaledSiLU); @@ -524,6 +530,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -543,6 +552,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc index 44747120cc2bc..92af2ee42a95b 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc @@ -15,6 +15,7 @@ ONNX_OPERATOR_KERNEL_EX( BranchwiseRMSNorm, kMSDomain, 1, kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", WebGpuSupportedFloatTypes()) + .TypeConstraint("F", DataTypeImpl::GetTensorType()) .TypeConstraint("M", WebGpuSupportedFloatTypes()), BranchwiseRMSNorm); @@ -61,7 +62,10 @@ Status BranchwiseRMSNormProgram::GenerateShaderCode(ShaderHelper& shader) const BranchwiseRMSNorm::BranchwiseRMSNorm(const OpKernelInfo& info) : WebGpuKernel(info), epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), - num_branches_(info.GetAttrOrDefault("num_branches", 0)) {} + num_branches_(info.GetAttrOrDefault("num_branches", 0)) { + ORT_ENFORCE(info.GetAttrOrDefault("use_onnx_arithmetic", 0) == 0, + "WebGPU BranchwiseRMSNorm does not support use_onnx_arithmetic=1"); +} Status BranchwiseRMSNorm::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { const auto* x = context.Input(0); diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index fecac2f17f523..761fc505b9530 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -3476,58 +3476,68 @@ constexpr const char* RotaryEmbedding_ver1_doc = R"DOC( RotaryEmbedding is the implementation of rotary positional embeddings (RoPE). The positions are represented as rotation matrices that are multiplied to query and key before the inner product of query and key is taken. )DOC"; +static OpSchema RotaryEmbeddingSchema() { + return OpSchema() + .SetDoc(RotaryEmbedding_ver1_doc) + .Attr("scale", + "Custom scale will be used if specified. Default value is 1.0", + AttributeProto::FLOAT, + OPTIONAL_VALUE) + .Attr("interleaved", + "Indicates whether the input has real and imaginary parts interleaved. " + "Default value is 0 (False), meaning the first half of the input consists of real values " + "and the second half consists of imaginary values.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("rotary_embedding_dim", + "Rotary embedding dimension. Default value is 0.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("num_heads", + "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("is_packed_batching", + "ragged batch inputs or not. Default value is 0", + AttributeProto::INT, + OPTIONAL_VALUE) + .Input(0, + "input", + "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", + "T") + .Input(1, + "position_ids", + "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", + "M") + .Input(2, + "cos_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Input(3, + "sin_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Output(0, + "output", + "tensor with same shape as input.", + "T") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") + .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateShapeFromInputToOutput(ctx, 0, 0); + }); +} + +ONNX_MS_OPERATOR_SET_SCHEMA(RotaryEmbedding, 1, RotaryEmbeddingSchema()); ONNX_MS_OPERATOR_SET_SCHEMA( - RotaryEmbedding, 1, - OpSchema() - .SetDoc(RotaryEmbedding_ver1_doc) - .Attr("scale", - "Custom scale will be used if specified. Default value is 1.0", - AttributeProto::FLOAT, - OPTIONAL_VALUE) - .Attr("interleaved", - "Indicates whether the input has real and imaginary parts interleaved. " - "Default value is 0 (False), meaning the first half of the input consists of real values " - "and the second half consists of imaginary values.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("rotary_embedding_dim", - "Rotary embedding dimension. Default value is 0.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("num_heads", - "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("is_packed_batching", - "ragged batch inputs or not. Default value is 0", - AttributeProto::INT, - OPTIONAL_VALUE) - .Input(0, - "input", - "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", - "T") - .Input(1, - "position_ids", - "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", - "M") - .Input(2, - "cos_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Input(3, - "sin_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Output(0, - "output", - "tensor with same shape as input.", - "T") - .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") - .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - propagateShapeFromInputToOutput(ctx, 0, 0); - })); + RotaryEmbedding, 2, + RotaryEmbeddingSchema() + .Attr("use_fp32_arithmetic", + "Widen input and cache values to float32, round each multiplication and addition in float32, " + "then convert the result to the input type. Values are 0 (legacy arithmetic) or 1. " + "Currently supported by CUDA only.", + AttributeProto::INT, static_cast(0))); constexpr const char* MRotaryEmbedding_ver1_doc = R"DOC( MRotaryEmbedding is the fused implementation of Multimodal Rotary Positional Embeddings (M-RoPE) used by the @@ -5754,31 +5764,90 @@ Applies RMS normalization independently to each branch. X may use grouped shape (..., C, H), or flattened shape (..., C * H) when num_branches is specified. The optional scale may have shape (C * H), (C, H), or (H). Arithmetic is performed in float32 and the result is converted to T. +With use_onnx_arithmetic=1, compute FP32 square, mean, epsilon addition, square +root, division, and scale multiplication with separate rounding boundaries. +CUDA reuses the standard fast last-axis ReduceMean summation order. This mode +does not substitute reciprocal square root or change the reduction tree. +Optional Y_fp32 losslessly widens the already T-rounded Y to float32; it does +not bypass the output rounding boundary and requires use_onnx_arithmetic=1. +WebGPU currently supports only the default mode without Y_fp32. )DOC"; +static OpSchema BranchwiseRMSNormSchema() { + return OpSchema() + .SetDoc(BranchwiseRMSNorm_ver1_doc) + .Attr("epsilon", "Epsilon added before reciprocal square root.", + AttributeProto::FLOAT, 1e-5f) + .Attr("num_branches", + "Number of branches for flattened input. Omit or set to zero for grouped input.", + AttributeProto::INT, static_cast(0)) + .Attr("use_onnx_arithmetic", + "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. " + "Supported values are 0 (default) and 1.", + AttributeProto::INT, static_cast(0)) + .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") + .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", + "M", OpSchema::Optional) + .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") + .Output(1, "Y_fp32", + "Optional float32 widening of the already T-rounded Y. " + "Requires use_onnx_arithmetic=1.", + "F", OpSchema::Optional) + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain input and output to floating-point tensors.") + .TypeConstraint("M", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain scale to floating-point tensors.") + .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + const auto start = getAttribute(ctx, "input_start", 0); + const auto width = getAttribute(ctx, "input_width", 0); + if (start < 0 || width < 0 || (width == 0 && start != 0)) { + fail_shape_inference("input_start and input_width must describe a nonnegative slice"); + } + if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 || + getAttribute(ctx, "num_branches", 0) != 0)) { + fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input"); + } + if (hasInputShape(ctx, 0)) { + auto shape = getInputShape(ctx, 0); + if (width != 0) { + if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two"); + auto* last = shape.mutable_dim(shape.dim_size() - 1); + if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) { + fail_shape_inference("input slice exceeds the last dimension"); + } + last->set_dim_value(width); + } + updateOutputShape(ctx, 0, shape); + } + if (ctx.getNumOutputs() > 1) { + if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) { + fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1"); + } + updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + if (hasInputShape(ctx, 0)) { + auto shape = getInputShape(ctx, 0); + if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width); + updateOutputShape(ctx, 1, shape); + } + } + }); +} + +ONNX_MS_OPERATOR_SET_SCHEMA(BranchwiseRMSNorm, 1, BranchwiseRMSNormSchema()); ONNX_MS_OPERATOR_SET_SCHEMA( - BranchwiseRMSNorm, 1, - OpSchema() - .SetDoc(BranchwiseRMSNorm_ver1_doc) - .Attr("epsilon", "Epsilon added before reciprocal square root.", - AttributeProto::FLOAT, 1e-5f) - .Attr("num_branches", - "Number of branches for flattened input. Omit or set to zero for grouped input.", + BranchwiseRMSNorm, 2, + BranchwiseRMSNormSchema() + .Attr("input_start", "CUDA-only starting offset in the last input dimension.", AttributeProto::INT, static_cast(0)) - .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") - .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", - "M", OpSchema::Optional) - .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") - .TypeConstraint("T", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain input and output to floating-point tensors.") - .TypeConstraint("M", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain scale to floating-point tensors.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - propagateShapeFromInputToOutput(ctx, 0, 0); - })); + .Attr("input_width", + "CUDA-only width selected from the last input dimension before normalization. " + "Zero retains the full input. Requires grouped input and use_onnx_arithmetic=1. " + "Outputs retain the leading dimensions and use this last dimension.", + AttributeProto::INT, static_cast(0))); constexpr const char* ScaledSiLU_ver1_doc = R"DOC( Computes SiLU after scaling, with explicit T rounding: diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index d68cdd392666b..93b771a425e07 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, 2, BranchwiseRMSNorm); 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); @@ -114,6 +115,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RemovePadding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RestorePadding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Rfft); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RotaryEmbedding); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 2, RotaryEmbedding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MRotaryEmbedding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GemmaRotaryEmbedding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SampleOp); @@ -225,6 +227,7 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); @@ -250,6 +253,7 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index 60ffd772c97f1..75024a420d3a9 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -275,7 +275,7 @@ __device__ void reduce_all( template __global__ void reduce_matrix_columns_kernel( const int num_rows, const int num_cols, const TIn* const input, TOut* const output, - TBuf* const block_reductions_buffer, int* const block_done_counts_buffer) { + TBuf* const block_reductions_buffer, int* const block_done_counts_buffer, const int input_row_stride) { const int num_blocks_in_grid_row = gridDim.x; const int row_id_in_grid = blockIdx.y; const int num_grid_rows = gridDim.y; @@ -283,7 +283,7 @@ __global__ void reduce_matrix_columns_kernel( // one row per iteration // row_id is int64_t to avoid int overflow in offset calculations for (int64_t row_id = row_id_in_grid; row_id < num_rows; row_id += num_grid_rows) { - const TIn* const row_data = input + row_id * num_cols; + const TIn* const row_data = input + row_id * input_row_stride; TOut* const row_output = output + row_id; TBuf* const row_block_reductions_buffer = block_reductions_buffer + row_id * num_blocks_in_grid_row; int* const row_block_done_counts_buffer = block_done_counts_buffer + row_id; @@ -296,8 +296,10 @@ __global__ void reduce_matrix_columns_kernel( template Status call_reduce_matrix_columns( - cudaStream_t stream, const TIn* input, TOut* output, const int num_rows, const int num_cols, void* buffer, size_t buffer_size) { + cudaStream_t stream, const TIn* input, TOut* output, const int num_rows, const int num_cols, + void* buffer, size_t buffer_size, int input_row_stride = 0) { ORT_ENFORCE(num_rows >= 0 && num_cols >= 0); + ORT_ENFORCE(input_row_stride == 0 || input_row_stride >= num_cols); using TBuf = AccumulationType_t; @@ -319,7 +321,8 @@ Status call_reduce_matrix_columns( const int shared_mem_size = sizeof(TBuf) * block_dim.x * block_dim.y / GPU_WARP_SIZE_HOST; reduce_matrix_columns_kernel <<>>( - num_rows, num_cols, input, output, block_reductions_buffer, block_done_counts_buffer); + num_rows, num_cols, input, output, block_reductions_buffer, block_done_counts_buffer, + input_row_stride == 0 ? num_cols : input_row_stride); return Status::OK(); } @@ -528,6 +531,26 @@ INSTANTIATE_REDUCE_MATRIX_COLUMNS(double); INSTANTIATE_REDUCE_MATRIX_COLUMNS(BFloat16); #undef INSTANTIATE_REDUCE_MATRIX_COLUMNS +struct RoundedSquare { + __device__ __forceinline__ float operator()(float value) const { + return __fmul_rn(value, value); + } +}; + +template +Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, + int m, int n, void* buffer, size_t buffer_size, int input_row_stride) { + return detail::call_reduce_matrix_columns( + stream, input, output, m, n, buffer, buffer_size, input_row_stride); +} + +#define INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(T) \ + template Status reduce_square_matrix_columns(cudaStream_t, const T*, float*, int, int, void*, size_t, int) +INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(float); +INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(half); +INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(BFloat16); +#undef INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS + namespace detail { constexpr int kMaxReduceRank = 16; diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index b772876db6395..9c877c0743b31 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h @@ -106,6 +106,11 @@ Status reduce_matrix_rows(cudaStream_t stream, const TIn* input, TOut* output, i template Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output, int m, int n, void* buffer, size_t buffer_size); +/** Sum FP32-rounded squares across columns, preserving reduce_matrix_columns' summation order. */ +template +Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, + int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0); + /** * Computes ReduceSum for an arbitrary set of axes. This is the general fallback * for axis layouts that cannot be flattened into one of the optimized matrix diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc index fbadf6da73678..52edede9208e0 100644 --- a/onnxruntime/core/session/environment.cc +++ b/onnxruntime/core/session/environment.cc @@ -296,7 +296,7 @@ Status Environment::Initialize(std::unique_ptr logging_ auto& domainToVersionRangeInstance = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance(); if (domainToVersionRangeInstance.Map().find(onnxruntime::kMSDomain) == domainToVersionRangeInstance.Map().end()) { // External shared providers may have already added kMSDomain - domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 1); + domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 2); } domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSExperimentalDomain, 1, 1); domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSNchwcDomain, 1, 1); diff --git a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc index 658c8695643e6..a6957e18bc194 100644 --- a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc +++ b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc @@ -2,6 +2,8 @@ // Licensed under the MIT License. #include +#include +#include #include #include "gtest/gtest.h" #include "contrib_ops/cpu/bert/rotary_embedding.h" @@ -21,6 +23,149 @@ enum class TensorType { }; } // anonymous namespace +class FloatRotaryChainTester : public OpTester { + public: + FloatRotaryChainTester(int type, int heads, int rotary_dim, bool interleaved) + : OpTester("RotaryEmbedding", 1, kMSDomain, false), + type_(type), + heads_(heads), + rotary_dim_(rotary_dim), + interleaved_(interleaved) {} + + void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, + std::vector>&) override { + auto value = [&](const std::string& name) { return &graph.GetOrCreateNodeArg(name, nullptr); }; + std::vector widened = inputs; + for (size_t index : {size_t{0}, size_t{2}, size_t{3}}) { + const auto name = "float_" + std::to_string(index); + widened[index] = value(name); + graph.AddNode(name, "Cast", "", {inputs[index]}, {widened[index]}) + .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); + } + auto* result = value("float_result"); + auto& rotary = graph.AddNode("rotary", "RotaryEmbedding", "", widened, {result}, nullptr, kMSDomain); + rotary.AddAttribute("num_heads", static_cast(heads_)); + rotary.AddAttribute("rotary_embedding_dim", static_cast(rotary_dim_)); + rotary.AddAttribute("interleaved", static_cast(interleaved_)); + graph.AddNode("rounded", "Cast", "", {result}, {outputs[0]}) + .AddAttribute("to", static_cast(type_)); + graph.SetOutputs(outputs); + } + + private: + void CreateModelToTest(const ModelOptions& options, Model*& model) override { + model = &BuildModel({{kOnnxDomain, 17}}, options); + ORT_THROW_IF_ERROR(model->MainGraph().Resolve()); + } + + int type_, heads_, rotary_dim_; + bool interleaved_; +}; + +template +void TestFloatRotaryArithmetic(int type) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA SM80 is required"; + } + for (int sequence : {1, 2, 9}) { + for (int width : {64, 256}) { + for (bool transposed : {false, true}) { + for (bool interleaved : {false, true}) { + SCOPED_TRACE("sequence=" + std::to_string(sequence) + " width=" + std::to_string(width) + + " transposed=" + std::to_string(transposed) + " interleaved=" + std::to_string(interleaved)); + constexpr int batch = 2, heads = 3; + const int rotary_dim = 64; + const int cache_rows = batch * sequence + 3; + std::vector shape = transposed ? std::vector{batch, heads, sequence, width} + : std::vector{batch, sequence, heads * width}; + std::vector input(batch * sequence * heads * width), cosine(cache_rows * rotary_dim / 2), + sine(cosine.size()); + const bool scalar_position = !transposed; + std::vector positions(scalar_position ? 1 : batch * sequence); + for (size_t i = 0; i < input.size(); ++i) input[i] = T(std::sin(i * 0.71f) * (1 + i % 17)); + for (size_t i = 0; i < cosine.size(); ++i) { + cosine[i] = T(std::cos(i * 0.13f)); + sine[i] = T(std::sin(i * 0.13f)); + } + for (size_t i = 0; i < positions.size(); ++i) positions[i] = scalar_position ? 1 : (i * 7) % cache_rows; + FloatRotaryChainTester reference(type, heads, rotary_dim, interleaved); + OpTester fused("RotaryEmbedding", 2, kMSDomain, false); + fused.AddAttribute("num_heads", heads); + fused.AddAttribute("rotary_embedding_dim", rotary_dim); + fused.AddAttribute("interleaved", interleaved); + fused.AddAttribute("use_fp32_arithmetic", 1); + std::vector expected; + for (OpTester* tester : {static_cast(&reference), &fused}) { + tester->AddInput("input", shape, input); + tester->AddInput("positions", scalar_position ? std::vector{1} : std::vector{batch, sequence}, positions); + tester->AddInput("cosine", {cache_rows, rotary_dim / 2}, cosine); + tester->AddInput("sine", {cache_rows, rotary_dim / 2}, sine); + tester->AddOutput("output", shape, std::vector(input.size())); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester->Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + auto fetches = tester->GetFetches(); + ASSERT_EQ(fetches.size(), 1); + if (expected.empty()) { + expected = std::move(fetches); + } else { + const auto& before = expected[0].Get(); + const auto& after = fetches[0].Get(); + ASSERT_EQ(before.SizeInBytes(), after.SizeInBytes()); + ASSERT_EQ(std::memcmp(before.DataRaw(), after.DataRaw(), before.SizeInBytes()), 0); + } + } + } + } + } + } +} + +TEST(RotaryEmbeddingTest, FloatArithmeticMatchesWidenRotateRoundFloat16) { + TestFloatRotaryArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); +} + +TEST(RotaryEmbeddingTest, FloatArithmeticMatchesWidenRotateRoundBFloat16) { + TestFloatRotaryArithmetic(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); +} + +TEST(RotaryEmbeddingTest, Version2LegacyDefaultAndInvalidMode) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA SM80 is required"; + } + for (int mode : {0, 2}) { + OpTester test("RotaryEmbedding", 2, kMSDomain); + if (mode != 0) test.AddAttribute("use_fp32_arithmetic", mode); + test.AddAttribute("num_heads", 1); + test.AddInput("input", {1, 1, 4}, ToFloat16({1, 2, 3, 4})); + test.AddInput("positions", {1}, {0}); + test.AddInput("cos", {1, 2}, ToFloat16({1, 1})); + test.AddInput("sin", {1, 2}, ToFloat16({0, 0})); + test.AddOutput("output", {1, 1, 4}, ToFloat16({1, 2, 3, 4})); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, + mode == 0 ? "" : "use_fp32_arithmetic must be 0 or 1", {}, nullptr, &providers); + } +} + +TEST(RotaryEmbeddingTest, Version2DoesNotFallBackToLegacyCpuArithmetic) { + OpTester test("RotaryEmbedding", 2, kMSDomain); + test.AddAttribute("use_fp32_arithmetic", 1); + test.AddAttribute("num_heads", 1); + test.AddInput("input", {1, 1, 4}, {1, 2, 3, 4}); + test.AddInput("positions", {1}, {0}); + test.AddInput("cos", {1, 2}, {1, 1}); + test.AddInput("sin", {1, 2}, {0, 0}); + test.AddOutput("output", {1, 1, 4}, {1, 2, 3, 4}); + std::vector> providers; + providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, "Could not find an implementation for RotaryEmbedding(2)", + {}, nullptr, &providers); +} + static void RunTest( const std::vector& input_data, const std::vector& position_ids, diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 92f803030ada4..207038805f624 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -5082,7 +5082,7 @@ void RegisterOrtOpSchemas() { auto& domainToVersionRangeInstance = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance(); if (domainToVersionRangeInstance.Map().find(onnxruntime::kMSDomain) == domainToVersionRangeInstance.Map().end()) { // External shared providers may have already added kMSDomain - domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 1); + domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 2); } domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSExperimentalDomain, 1, 1); domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSNchwcDomain, 1, 1); From 80988834676fc52705fe091008549c67a5c0ab6a Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 19:25:53 +0000 Subject: [PATCH 2/8] Cover exact and packed RMS normalization on standalone main branch --- .../contrib_ops/hyper_connection_op_test.cc | 252 ++++++++++++++++++ 1 file changed, 252 insertions(+) diff --git a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc index f8f590ec32598..622116282406c 100644 --- a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc +++ b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc @@ -2,7 +2,9 @@ // Licensed under the MIT License. #include +#include #include +#include #include #include @@ -10,6 +12,7 @@ #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" +#include "core/framework/session_options.h" namespace onnxruntime::test { namespace { @@ -127,6 +130,255 @@ TEST(HyperConnectionOpsTest, BranchwiseRMSNormMixedScaleType) { }); } +class StandardRMSChainTester : public OpTester { + public: + StandardRMSChainTester(int output_type, float epsilon, int64_t start = 0, int64_t width = 0) + : OpTester("StandardRMSChain", 17, kOnnxDomain, false), + output_type_(output_type), + epsilon_(epsilon), + start_(start), + width_(width) {} + + void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, + std::vector>&) override { + auto value = [&](const char* name) { return &graph.GetOrCreateNodeArg(name, nullptr); }; + NodeArg* source = inputs[0]; + if (width_ != 0) { + auto constant = [&](const char* name, int64_t data) { + ONNX_NAMESPACE::TensorProto tensor; + tensor.set_name(name); + tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + tensor.add_dims(1); + tensor.add_int64_data(data); + graph.AddInitializedTensor(tensor); + return value(name); + }; + source = value("selected"); + graph.AddNode("slice", "Slice", "", + {inputs[0], constant("starts", start_), constant("ends", start_ + width_), + constant("axes", -1)}, + {source}); + } + auto* x = value("float_x"); + graph.AddNode("cast", "Cast", "", {source}, {x}) + .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); + auto* square = value("square"); + graph.AddNode("square", "Mul", "", {x, x}, {square}); + auto* mean = value("mean"); + auto& reduce = graph.AddNode("mean", "ReduceMean", "", {square}, {mean}); + reduce.AddAttribute("axes", std::vector{-1}); + reduce.AddAttribute("keepdims", int64_t{1}); + ONNX_NAMESPACE::TensorProto epsilon; + epsilon.set_name("epsilon"); + epsilon.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + epsilon.add_float_data(epsilon_); + graph.AddInitializedTensor(epsilon); + auto* adjusted = value("adjusted"); + graph.AddNode("epsilon", "Add", "", {mean, value("epsilon")}, {adjusted}); + auto* root = value("root"); + graph.AddNode("root", "Sqrt", "", {adjusted}, {root}); + auto* normalized = value("normalized"); + graph.AddNode("divide", "Div", "", {x, root}, {normalized}); + auto* scaled = value("scaled"); + graph.AddNode("scale", "Mul", "", {normalized, inputs[1]}, {scaled}); + graph.AddNode("output", "Cast", "", {scaled}, {outputs[0]}) + .AddAttribute("to", static_cast(output_type_)); + if (outputs.size() > 1) { + graph.AddNode("widen", "Cast", "", {outputs[0]}, {outputs[1]}) + .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); + } + graph.SetOutputs(outputs); + } + + private: + int output_type_; + float epsilon_; + int64_t start_; + int64_t width_; +}; + +template +void TestStandardRMSArithmetic(int type, bool packed = false) { + if (!DefaultCudaExecutionProvider()) { + GTEST_SKIP() << "CUDA is required for exact standard RMS arithmetic"; + } + for (int width : {1, 31, 128, 256, 513, 1024, 4096}) { + for (int rows : (packed ? std::vector{1, 2, 9, 48} : std::vector{1, 2, 9})) { + for (int start : (packed ? std::vector{0, 3} : std::vector{0})) { + SCOPED_TRACE("width=" + std::to_string(width) + " rows=" + std::to_string(rows)); + const int stride = packed ? width * 2 + 7 : width; + std::vector input(rows * stride); + std::vector scale(width); + for (int i = 0; i < rows * stride; ++i) { + input[i] = T(std::sin(static_cast(i) * 0.71f) * (1 + i % 17)); + } + for (int i = 0; i < width; ++i) scale[i] = 0.5f + static_cast(i % 23) / 17.0f; + constexpr float epsilon = 1e-6f; + StandardRMSChainTester reference(type, epsilon, start, packed ? width : 0); + OpTester fused("BranchwiseRMSNorm", packed ? 2 : 1, kMSDomain, false); + if (packed) { + fused.AddAttribute("input_start", start); + fused.AddAttribute("input_width", width); + } + fused.AddAttribute("use_onnx_arithmetic", 1); + fused.AddAttribute("epsilon", epsilon); + std::vector expected; + for (OpTester* tester : {static_cast(&reference), &fused}) { + const std::vector input_shape = rows == 48 ? std::vector{1, 2, 24, stride} + : std::vector{1, rows, stride}; + auto output_shape = input_shape; + output_shape.back() = width; + tester->AddInput("X", input_shape, input); + tester->AddInput("scale", {width}, scale); + tester->AddOutput("Y", output_shape, std::vector(rows * width)); + tester->AddOutput("Y_fp32", output_shape, std::vector(rows * width)); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester->Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + auto fetches = tester->GetFetches(); + ASSERT_EQ(fetches.size(), 2); + if (expected.empty()) { + expected = std::move(fetches); + } else { + for (size_t output = 0; output < fetches.size(); ++output) { + const auto& before = expected[output].Get(); + const auto& after = fetches[output].Get(); + ASSERT_EQ(before.SizeInBytes(), after.SizeInBytes()); + EXPECT_EQ(std::memcmp(before.DataRaw(), after.DataRaw(), before.SizeInBytes()), 0) + << "output=" << output; + } + } + } + } + } + } +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardFloatExact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardFloat16Exact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardBFloat16Exact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedFloatExact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, true); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedFloat16Exact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, true); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedBFloat16Exact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16, true); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedRejectsInvalidSlice) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for packed RMS"; + for (const auto& [start, width, mode, branches, message] : + std::vector>{ + {-1, 2, 1, 0, "nonnegative"}, + {0, -1, 1, 0, "nonnegative"}, + {1, 0, 1, 0, "nonnegative"}, + {0, 5, 1, 0, "exceeds"}, + {3, 2, 1, 0, "exceeds"}, + {0, 2, 0, 0, "requires"}, + {0, 2, 1, 2, "requires"}}) { + OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); + tester.AddAttribute("input_start", start); + tester.AddAttribute("input_width", width); + tester.AddAttribute("use_onnx_arithmetic", mode); + tester.AddAttribute("num_branches", branches); + tester.AddInput("X", {1, 1, 4}, {1, 2, 3, 4}); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {1, 1, 2}, {0, 0}); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, message, {}, nullptr, &providers); + } +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedEmptyWithoutScale) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for packed RMS"; + OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 1); + tester.AddAttribute("input_start", 3); + tester.AddAttribute("input_width", 2); + tester.AddInput("X", {0, 3, 8}, {}); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {0, 3, 2}, {}); + tester.AddOutput("Y_fp32", {0, 3, 2}, {}); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedWithoutScale) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for packed RMS"; + OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 1); + tester.AddAttribute("input_start", 1); + tester.AddAttribute("input_width", 2); + tester.AddAttribute("epsilon", 0.0f); + tester.AddInput("X", {1, 2, 4}, {99, 3, 4, 99, 99, 0, 2, 99}); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {1, 2, 2}, + {3.0f / std::sqrt(12.5f), 4.0f / std::sqrt(12.5f), 0, 2.0f / std::sqrt(2.0f)}); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedDoesNotFallBackToCpu) { + OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 1); + tester.AddAttribute("input_width", 2); + tester.AddInput("X", {1, 1, 4}, {1, 2, 3, 4}); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {1, 1, 2}, {0, 0}); + std::vector> providers; + providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, "Could not find an implementation for BranchwiseRMSNorm(2)", + {}, nullptr, &providers); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardCpuBranchScale) { + OpTester tester("BranchwiseRMSNorm", 1, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 1); + tester.AddAttribute("num_branches", 2); + tester.AddAttribute("epsilon", 0.0f); + tester.AddInput("X", {1, 4}, {3.0f, 4.0f, 0.0f, 2.0f}); + tester.AddInput("scale", {2, 2}, {2.0f, 0.5f, 0.25f, 3.0f}); + tester.AddOutput("Y", {1, 4}, + {3.0f / std::sqrt(12.5f) * 2.0f, 4.0f / std::sqrt(12.5f) * 0.5f, + 0.0f, 2.0f / std::sqrt(2.0f) * 3.0f}); + std::vector> providers; + providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormRejectsInvalidArithmeticMode) { + for (auto& provider : AvailableProviders()) { + if (provider->Type() == kWebGpuExecutionProvider) continue; + OpTester tester("BranchwiseRMSNorm", 1, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 2); + tester.AddInput("X", {1, 2}, {1.0f, 2.0f}); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {1, 2}, {0.0f, 0.0f}); + std::vector> providers; + providers.push_back(std::move(provider)); + tester.Run(OpTester::ExpectResult::kExpectFailure, "use_onnx_arithmetic must be 0 or 1", + {}, nullptr, &providers); + } +} + TEST(HyperConnectionOpsTest, ScaledSiLU) { RunOnAvailableProviders("ScaledSiLU", [&](OpTester& tester) { tester.AddAttribute("alpha", 0.5f); From 17662fc00406e9058cca4a57b1907fa52b1acd76 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 20:47:47 +0000 Subject: [PATCH 3/8] Fix normalization determinism and preserve exact rotary fusion semantics --- docs/ContribOperators.md | 65 ++++------ docs/OperatorKernels.md | 8 +- .../cuda/bert/branchwise_rms_norm.cc | 6 +- .../cuda/bert/branchwise_rms_norm_impl.cu | 8 +- .../cuda/bert/branchwise_rms_norm_impl.h | 3 +- .../core/graph/contrib_ops/bert_defs.cc | 6 +- .../optimizer/group_query_attention_fusion.cc | 3 + .../contrib_ops/hyper_connection_op_test.cc | 116 ++++++++++++++++++ .../graph_transform_test_layernorm.cc | 30 ++++- 9 files changed, 194 insertions(+), 51 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index d1594969658f2..0ad97f90ebea8 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -896,37 +896,33 @@ This version of the operator has been available since version 1 of the 'com.micr performed in float32 and the result is converted to T. With use_onnx_arithmetic=1, compute FP32 square, mean, epsilon addition, square root, division, and scale multiplication with separate rounding boundaries. - CUDA reuses the standard fast ReduceMean summation order, including its scalar-result - path; it does not substitute reciprocal square root or change the reduction tree. - For multiple rows, FP32-rounded squaring is performed during the reduction load rather - than materializing a separate FP32 square tensor. - Optional Y_fp32 losslessly widens the already T-rounded Y and requires - use_onnx_arithmetic=1. WebGPU supports only the default mode without Y_fp32. - CUDA-only version 2 can select a contiguous last-axis interval directly from - packed rows using input_start and input_width. It retains the same reduction - order over the selected width and avoids materializing a Slice tensor. + CUDA reuses the standard fast last-axis ReduceMean summation order by default. + Deterministic CUDA sessions instead use a fixed-order reduction, which may + round differently from the default reduction or deterministic ReduceMean. + This mode does not substitute reciprocal square root. + Optional Y_fp32 losslessly widens the already T-rounded Y to float32; it does + not bypass the output rounding boundary and requires use_onnx_arithmetic=1. + WebGPU currently supports only the default mode without Y_fp32. #### Version -Version 1 retains the original input layout. Version 2 adds CUDA-only packed-input -selection; CPU and WebGPU implement version 1 only. +This version of the operator has been available since version 2 of the 'com.microsoft' operator set. + +Other versions of this operator: com.microsoft.BranchwiseRMSNorm-1 #### Attributes
-
epsilon : float
+
epsilon : float (default is 1e-05)
Epsilon added before reciprocal square root.
-
input_start : int (default is 0, version 2)
-
Nonnegative starting offset in the last input dimension. A nonzero offset requires input_width.
-
input_width : int (default is 0, version 2)
-
Width selected before normalization. Zero retains the full input. A positive width -requires use_onnx_arithmetic=1 and num_branches=0, and the selected interval must fit -within the last input dimension.
+
input_start : int (default is 0)
+
CUDA-only starting offset in the last input dimension.
+
input_width : int (default is 0)
+
CUDA-only width selected from the last input dimension before normalization. Zero retains the full input. Requires grouped input and use_onnx_arithmetic=1. Outputs retain the leading dimensions and use this last dimension.
num_branches : int (default is 0)
Number of branches for flattened input. Omit or set to zero for grouped input.
use_onnx_arithmetic : int (default is 0)
-
Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. -Supported values are 0 (default) and 1.
+
Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. Supported values are 0 (default) and 1.
#### Inputs (1 - 2) @@ -938,25 +934,24 @@ Supported values are 0 (default) and 1.
Optional scale with shape (C * H), (C, H), or (H).
-#### Outputs +#### Outputs (1 - 2)
Y : T
-
RMS-normalized output with the same shape as X, except that version 2 with -positive input_width replaces the last dimension with input_width.
+
RMS-normalized output with the same shape as X.
Y_fp32 (optional) : F
-
Float32 widening of the already T-rounded Y. Requires use_onnx_arithmetic=1.
+
Optional float32 widening of the already T-rounded Y. Requires use_onnx_arithmetic=1.
#### Type Constraints
-
F : tensor(float)
-
Widened output is float32.
T : tensor(float), tensor(float16), tensor(bfloat16)
Constrain input and output to floating-point tensors.
M : tensor(float), tensor(float16), tensor(bfloat16)
Constrain scale to floating-point tensors.
+
F : tensor(float)
+
Widened output is float32.
@@ -7199,18 +7194,12 @@ This version of the operator has been available since version 1 of the 'com.micr RotaryEmbedding is the implementation of rotary positional embeddings (RoPE). The positions are represented as rotation matrices that are multiplied to query and key before the inner product of query and key is taken. - Version 2 adds optional FP32 arithmetic with activation-type storage. With - use_fp32_arithmetic=1, input and cache values are widened inside the CUDA kernel, - each multiplication and addition is rounded separately in float32, and only the - final result is converted back to the input type. This preserves a - Cast-to-float -> float RotaryEmbedding -> Cast-to-input-type chain. - Version 1 and the default arithmetic are unchanged. #### Version -The latest version is 2 of the 'com.microsoft' operator set. Version 2 is currently -implemented by CUDA only; other providers retain version 1 and cannot silently -execute version 2 with legacy low-precision arithmetic. +This version of the operator has been available since version 2 of the 'com.microsoft' operator set. + +Other versions of this operator: com.microsoft.RotaryEmbedding-1 #### Attributes @@ -7225,10 +7214,8 @@ execute version 2 with legacy low-precision arithmetic.
Rotary embedding dimension. Default value is 0.
scale : float
Custom scale will be used if specified. Default value is 1.0
-
use_fp32_arithmetic : int (default is 0; version 2 only)
-
Widen input and cache values, perform separately rounded float32 multiplication -and addition, then round the output once to the input type. Supported values are -0 (legacy arithmetic) and 1.
+
use_fp32_arithmetic : int (default is 0)
+
Widen input and cache values to float32, round each multiplication and addition in float32, then convert the result to the input type. Values are 0 (legacy arithmetic) or 1. Currently supported by CUDA only.
#### Inputs diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 669e3083f6b15..46971132fefeb 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -575,7 +575,7 @@ The **OpSet Version** column uses the following notation: |BeamSearch|*in* input_ids:**F**
*in* max_length:**I**
*in* min_length:**I**
*in* num_beams:**I**
*in* num_return_sequences:**I**
*in* length_penalty:**T**
*in* repetition_penalty:**T**
*in* vocab_mask:**M**
*in* prefix_vocab_mask:**M**
*in* attention_mask:**I**
*in* decoder_input_ids:**I**
*in* logits_processor:**I**
*out* sequences:**I**
*out* sequences_scores:**T**
*out* scores:**T**|1+|**T** = tensor(float)| |BiasGelu|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float)| |BifurcationDetector|*in* src_tokens:**T**
*in* cur_tokens:**T**
*in* prev_suffix_match_idx:**T**
*in* pred_tokens:**T**
*out* tokens:**T**
*out* suffix_match_idx:**T**|1+|**T** = tensor(int64)| -|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**
*out* Y_fp32:**F**|1+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |CDist|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(double), tensor(float)| |CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**|1+|**T** = tensor(float)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| @@ -1088,7 +1088,8 @@ The **OpSet Version** column uses the following notation: |BiasSplitGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |BitmaskBiasDropout|*in* data:**T**
*in* bias:**T**
*in* residual:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| |BitmaskDropout|*in* data:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| -|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**
*out* Y_fp32:**F**|2+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|||1+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |ComplexMul|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ComplexMulConj|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| @@ -1153,7 +1154,8 @@ The **OpSet Version** column uses the following notation: |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)| |Rfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| -|RotaryEmbedding|*in* input:**T**
*in* position_ids:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**|1+|**M** = tensor(int64)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|RotaryEmbedding|*in* input:**T**
*in* position_ids:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**|2+|**M** = tensor(int64)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|||1+|**M** = tensor(int64)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |Sampling|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* repetition_penalty:**T**
*in* vocab_mask:**I**
*in* prefix_vocab_mask:**I**
*in* attention_mask:**I**
*in* presence_mask:**I**
*in* seed:**I**
*out* sequences:**I**
*out* filtered_logits:**T**|1+|**T** = tensor(float), tensor(float16)| |ScaledSiLU|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |SkipGroupNorm|*in* X:**T**
*in* gamma:**M**
*in* beta:**M**
*in* skip:**T**
*in* bias:**T**
*out* Y:**T**
*out* S:**T**|1+|**T** = tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc index 1c91e9619364e..742294340a02e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc @@ -91,7 +91,9 @@ Status BranchwiseRMSNorm::ComputeInternal(OpKernelContext* context) const { "ONNX arithmetic requires the row count to fit in int32"); const auto buffer_size = compute_reduce_matrix_columns_buffer_size( static_cast(groups), static_cast(params.hidden)); - auto squares = GetScratchBuffer(groups == 1 ? output_shape.Size() : 0, GetComputeStream(context)); + const bool deterministic_compute = context->GetUseDeterministicCompute(); + auto squares = GetScratchBuffer(groups == 1 && !deterministic_compute ? output_shape.Size() : 0, + GetComputeStream(context)); auto sums = GetScratchBuffer(groups, GetComputeStream(context)); auto buffer = GetScratchBuffer(buffer_size, GetComputeStream(context)); return LaunchStandardBranchwiseRMSNorm( @@ -101,7 +103,7 @@ Status BranchwiseRMSNorm::ComputeInternal(OpKernelContext* context) const { scale == nullptr || scale->Shape().Size() == params.hidden, epsilon_, widened == nullptr ? nullptr : widened->MutableData(), squares.get(), sums.get(), buffer.get(), buffer_size, - static_cast(input_width_ == 0 ? params.hidden : input_stride)); + static_cast(input_width_ == 0 ? params.hidden : input_stride), deterministic_compute); } return LaunchMixedScaleBranchwiseRMSNorm( Stream(context), x_data, scale == nullptr ? nullptr : scale->DataRaw(), diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu index 240b057b827e0..ebcc0a3dee1ca 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu @@ -112,12 +112,14 @@ Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const vo int scale_type, T* y, int groups, int branches, int hidden, bool shared_scale, float epsilon, float* widened, float* squares, float* sums, void* reduction_buffer, - size_t reduction_buffer_size, int input_row_stride) { + size_t reduction_buffer_size, int input_row_stride, + bool deterministic_compute) { const int64_t count = static_cast(groups) * hidden; const int64_t blocks = (count + kThreads - 1) / kThreads; ORT_RETURN_IF_NOT(blocks <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); // Standard ReduceMean removes singleton dimensions and prefers row reduction for a scalar result. - if (groups == 1) { + // Row reduction uses floating-point atomics; column reduction combines indexed partials in a fixed order. + if (groups == 1 && !deterministic_compute) { BranchwiseSquareKernel<<(blocks), kThreads, 0, stream>>>(x, squares, count); ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); ORT_RETURN_IF_ERROR((onnxruntime::cuda::reduce_matrix_rows( @@ -170,7 +172,7 @@ INSTANTIATE(__nv_bfloat16) #define INSTANTIATE_STANDARD(T) \ template Status LaunchStandardBranchwiseRMSNorm(cudaStream_t, const T*, const void*, int, T*, \ - int, int, int, bool, float, float*, float*, float*, void*, size_t, int); + int, int, int, bool, float, float*, float*, float*, void*, size_t, int, bool); INSTANTIATE_STANDARD(float) INSTANTIATE_STANDARD(half) INSTANTIATE_STANDARD(__nv_bfloat16) diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h index a037130d825a2..6a80862d62657 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h @@ -19,6 +19,7 @@ Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const vo int scale_type, T* y, int groups, int branches, int hidden, bool shared_scale, float epsilon, float* widened, float* squares, float* sums, void* reduction_buffer, - size_t reduction_buffer_size, int input_row_stride); + size_t reduction_buffer_size, int input_row_stride, + bool deterministic_compute); } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 761fc505b9530..d233019398d82 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -5766,8 +5766,10 @@ The optional scale may have shape (C * H), (C, H), or (H). Arithmetic is performed in float32 and the result is converted to T. With use_onnx_arithmetic=1, compute FP32 square, mean, epsilon addition, square root, division, and scale multiplication with separate rounding boundaries. -CUDA reuses the standard fast last-axis ReduceMean summation order. This mode -does not substitute reciprocal square root or change the reduction tree. +CUDA reuses the standard fast last-axis ReduceMean summation order by default. +Deterministic CUDA sessions instead use a fixed-order reduction, which may +round differently from the default reduction or deterministic ReduceMean. +This mode does not substitute reciprocal square root. Optional Y_fp32 losslessly widens the already T-rounded Y to float32; it does not bypass the output rounding boundary and requires use_onnx_arithmetic=1. WebGPU currently supports only the default mode without Y_fp32. diff --git a/onnxruntime/core/optimizer/group_query_attention_fusion.cc b/onnxruntime/core/optimizer/group_query_attention_fusion.cc index 6506dfae7898d..d1903c297f210 100644 --- a/onnxruntime/core/optimizer/group_query_attention_fusion.cc +++ b/onnxruntime/core/optimizer/group_query_attention_fusion.cc @@ -352,6 +352,9 @@ static bool TryGetRotaryEmbeddingArgs(Node& rotary_node, RotaryEmbeddingArgs& ar if (rotary_node.Domain() != kMSDomain) { return false; } + if (GetIntAttributeOrDefault(rotary_node, "use_fp32_arithmetic", 0) != 0) { + return false; + } args.interleaved = GetIntAttributeOrDefault(rotary_node, "interleaved", 0); args.rotary_embedding_dim = GetIntAttributeOrDefault(rotary_node, "rotary_embedding_dim", 0); diff --git a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc index 622116282406c..f3efaa137d062 100644 --- a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc +++ b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc @@ -280,6 +280,122 @@ TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedBFloat16Exact) { TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16, true); } +template +void TestStandardRMSBranchScales(int type) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for exact branch scales"; + constexpr int outer = 2; + constexpr int branches = 3; + constexpr float epsilon = 1e-6f; + for (int width : {31, 128, 513}) { + for (bool flattened_scale : {false, true}) { + SCOPED_TRACE(width); + SCOPED_TRACE(flattened_scale); + std::vector input(outer * branches * width); + std::vector scale(branches * width); + std::vector float_scale(scale.size()); + for (size_t i = 0; i < input.size(); ++i) input[i] = T(std::sin(static_cast(i) * 0.71f) * 9.0f); + for (int branch = 0; branch < branches; ++branch) { + for (int h = 0; h < width; ++h) { + const int index = branch * width + h; + scale[index] = ScaleT((branch + 1) * 0.7f + (h % 17) * 0.013f); + if constexpr (std::is_same_v) { + float_scale[index] = scale[index]; + } else { + float_scale[index] = scale[index].ToFloat(); + } + } + } + const std::vector shape{outer, branches, width}; + StandardRMSChainTester reference(type, epsilon); + reference.AddInput("X", shape, input); + reference.AddInput("scale", {branches, width}, float_scale); + reference.AddOutput("Y", shape, std::vector(input.size())); + reference.AddOutput("Y_fp32", shape, std::vector(input.size())); + OpTester fused("BranchwiseRMSNorm", 2, kMSDomain, false); + fused.AddAttribute("use_onnx_arithmetic", 1); + fused.AddAttribute("epsilon", epsilon); + fused.AddInput("X", shape, input); + fused.AddInput("scale", flattened_scale ? std::vector{branches * width} : std::vector{branches, width}, + scale); + fused.AddOutput("Y", shape, std::vector(input.size())); + fused.AddOutput("Y_fp32", shape, std::vector(input.size())); + std::vector expected; + for (OpTester* tester : {static_cast(&reference), &fused}) { + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester->Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + const auto& fetches = tester->GetFetches(); + ASSERT_EQ(fetches.size(), 2u); + if (expected.empty()) { + expected = fetches; + } else { + for (size_t output = 0; output < fetches.size(); ++output) { + const auto& before = expected[output].Get(); + const auto& after = fetches[output].Get(); + ASSERT_EQ(before.SizeInBytes(), after.SizeInBytes()); + EXPECT_EQ(std::memcmp(before.DataRaw(), after.DataRaw(), before.SizeInBytes()), 0) << output; + } + } + } + } + } +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormCudaExactBranchScales) { + TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); + TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); + TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); + TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); + TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); + TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormWideSingleRowDeterministic) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for deterministic RMS"; + constexpr int width = 8192; + constexpr float epsilon = 1e-6f; + std::vector input(width); + double square_sum = 0.0; + for (int i = 0; i < width; ++i) { + input[i] = std::sin(static_cast(i) * 0.71f) * (1 + i % 17); + const float square = input[i] * input[i]; + square_sum += square; + } + const float rms = std::sqrt(static_cast(square_sum / width) + epsilon); + std::vector expected(width); + for (int i = 0; i < width; ++i) expected[i] = input[i] / rms; + std::vector first; + for (int run = 0; run < 10; ++run) { + OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 1); + tester.AddAttribute("epsilon", epsilon); + tester.AddInput("X", {1, 1, width}, input); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {1, 1, width}, expected, false, 1e-5f, 1e-5f); + tester.AddOutput("Y_fp32", {1, 1, width}, expected, false, 1e-5f, 1e-5f); + SessionOptions options; + options.use_deterministic_compute = true; + options.graph_optimization_level = TransformerLevel::Default; + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + const auto& fetches = tester.GetFetches(); + ASSERT_EQ(fetches.size(), 2u); + const float* actual = fetches[0].Get().Data(); + if (first.empty()) { + first.assign(actual, actual + width); + } else { + EXPECT_EQ(std::memcmp(first.data(), actual, width * sizeof(float)), 0); + } + EXPECT_EQ(std::memcmp(actual, fetches[1].Get().Data(), width * sizeof(float)), 0); + } +} + TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedRejectsInvalidSlice) { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for packed RMS"; for (const auto& [start, width, mode, branches, message] : diff --git a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc index 6696790410603..14bc41c6da47b 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc +++ b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc @@ -1237,7 +1237,9 @@ static void BuildRotaryEmbeddingGQAFusionGraph(ModelTestBuilder& builder, int64_t q_interleaved = 0, int64_t k_interleaved = 0, int64_t rotary_embedding_dim = 0, - int64_t gqa_num_heads = 2) { + int64_t gqa_num_heads = 2, + bool q_fp32_arithmetic = false, + bool k_fp32_arithmetic = false) { constexpr int64_t batch_size = 1; constexpr int64_t sequence_length = 2; constexpr int64_t input_hidden_size = 8; @@ -1307,6 +1309,10 @@ static void BuildRotaryEmbeddingGQAFusionGraph(ModelTestBuilder& builder, k_rotary.AddAttribute("num_heads", kv_num_heads); k_rotary.AddAttribute("interleaved", k_interleaved); k_rotary.AddAttribute("rotary_embedding_dim", rotary_embedding_dim); + if (q_fp32_arithmetic || k_fp32_arithmetic) { + q_rotary.AddAttribute("use_fp32_arithmetic", static_cast(q_fp32_arithmetic)); + k_rotary.AddAttribute("use_fp32_arithmetic", static_cast(k_fp32_arithmetic)); + } NodeArg* past_key = builder.MakeInput({{batch_size, kv_num_heads, max_sequence_length, head_size}}); @@ -1718,6 +1724,28 @@ TEST_F(GraphTransformationTests, GroupQueryAttentionFusionMsRotaryEmbeddingForwa check_fused_graph)); } +TEST_F(GraphTransformationTests, GroupQueryAttentionFusionPreservesExactRotaryArithmetic) { + for (const auto& [q_exact, k_exact] : + {std::pair{false, false}, std::pair{true, false}, std::pair{false, true}, std::pair{true, true}}) { + Model model("ExactRotaryGQAFusion", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 23}, {kMSDomain, 2}}, {}, *logger_); + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + BuildRotaryEmbeddingGQAFusionGraph(builder, RotaryEmbeddingDomain::kMS, true, 1, 1, 0, 2, + q_exact, k_exact); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + GraphTransformerManager transformers{3}; + ASSERT_STATUS_OK(transformers.Register(std::make_unique(), TransformerLevel::Level2)); + ASSERT_STATUS_OK(transformers.ApplyTransformers(graph, TransformerLevel::Level2, *logger_)); + if (q_exact || k_exact) { + ASSERT_STATUS_OK(CheckMsRotaryEmbeddingGQANotFused(graph)); + } else { + ASSERT_STATUS_OK(CheckMsRotaryEmbeddingGQAFused(graph, 1)); + } + } +} + TEST_F(GraphTransformationTests, GroupQueryAttentionFusionSkipsMismatchedProjectionSizesTest) { auto build_test_case = [](ModelTestBuilder& builder) { BuildRotaryEmbeddingGQAFusionGraph(builder, RotaryEmbeddingDomain::kMS, true, 0, 0, 0, 3); From 5c027174e06e082136ac43c0cfd2ec325f8771b7 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 22 Oct 2026 05:43:25 +0000 Subject: [PATCH 4/8] Add fused BF16 RMS normalization and rotary embedding Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/contrib_ops/cuda/rms_norm_rotary.md | 62 +++++++ .../cpu/bert/branchwise_rms_norm.cc | 4 + .../contrib_ops/cpu/bert/rotary_embedding.cc | 5 + .../cuda/bert/branchwise_rms_norm.cc | 15 +- .../cuda/bert/branchwise_rms_norm_impl.cu | 13 +- .../cuda/bert/branchwise_rms_norm_impl.h | 2 +- .../contrib_ops/cuda/bert/rms_norm_rotary.cc | 115 ++++++++++++ .../contrib_ops/cuda/bert/rms_norm_rotary.cu | 88 +++++++++ .../contrib_ops/cuda/bert/rms_norm_rotary.h | 27 +++ .../contrib_ops/cuda/bert/rotary_embedding.cc | 30 ++- .../cuda/bert/rotary_embedding_impl.cu | 4 + .../contrib_ops/cuda/cuda_contrib_kernels.cc | 2 + .../webgpu/bert/branchwise_rms_norm.cc | 5 + .../webgpu/bert/rotary_embedding.cc | 8 +- .../core/graph/contrib_ops/bert_defs.cc | 87 +++++++++ onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 + .../test/contrib_ops/rms_norm_rotary_test.cc | 173 ++++++++++++++++++ 17 files changed, 626 insertions(+), 16 deletions(-) create mode 100644 docs/contrib_ops/cuda/rms_norm_rotary.md create mode 100644 onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc create mode 100644 onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu create mode 100644 onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h create mode 100644 onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc 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/cpu/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc index 23e8b6b02b354..5b093a28f72aa 100644 --- a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc @@ -67,6 +67,10 @@ Status BranchwiseRMSNorm::PrePack(const Tensor& tensor, int input_idx, Alloca template Status BranchwiseRMSNorm::Compute(OpKernelContext* context) const { + const auto* numerics = context->Input(2); + ORT_RETURN_IF_NOT(numerics == nullptr || + (numerics->Shape().Size() == 1 && *numerics->Data() == 0), + "BranchwiseRMSNorm: nonzero prefill_numerics requires CUDA"); const auto* x = context->Input(0); const auto* scale = context->Input(1); hyper_connection::StreamShape params; diff --git a/onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc b/onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc index 682c254840e91..c9b9e5a8626f7 100644 --- a/onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc +++ b/onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc @@ -184,6 +184,11 @@ template Status RunRotaryEmbedding(concurrency::ThreadPool* tp, Rotar template Status RotaryEmbedding::Compute(OpKernelContext* context) const { + const auto* numerics = context->Input(4); + ORT_RETURN_IF_NOT(numerics == nullptr || + (numerics->IsDataType() && numerics->Shape().Size() == 1 && + numerics->Data()[0] == 0), + "RotaryEmbedding: nonzero prefill_numerics requires CUDA"); const Tensor* input = context->Input(0); const Tensor* position_ids = context->Input(1); const Tensor* cos_cache = context->Input(2); diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc index 742294340a02e..bc6b356aba41d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc @@ -21,6 +21,8 @@ using namespace onnxruntime::cuda; (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ .TypeConstraint("F", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ + .InputMemoryType(OrtMemTypeCPUInput, 2) \ .TypeConstraint("M", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}), \ @@ -56,6 +58,17 @@ Status BranchwiseRMSNorm::ComputeInternal(OpKernelContext* context) const { using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; const auto* x = context->Input(0); const auto* scale = context->Input(1); + const auto* numerical_mode = context->Input(2); + bool dense_prefill = false; + if (numerical_mode != nullptr) { + ORT_RETURN_IF_NOT(numerical_mode->Shape().Size() == 1, + "prefill_numerics must contain one int32 value"); + const int32_t mode = *numerical_mode->Data(); + ORT_RETURN_IF_NOT(mode == 0 || mode == 1, "prefill_numerics must be 0 or 1"); + dense_prefill = mode == 1; + ORT_RETURN_IF_NOT(!dense_prefill || use_onnx_arithmetic_, + "Dense prefill numerical mode requires use_onnx_arithmetic=1"); + } auto output_dims = x->Shape().AsShapeVector(); ORT_RETURN_IF_NOT(!output_dims.empty(), "BranchwiseRMSNorm requires a non-scalar input"); const int64_t input_stride = output_dims.back(); @@ -103,7 +116,7 @@ Status BranchwiseRMSNorm::ComputeInternal(OpKernelContext* context) const { scale == nullptr || scale->Shape().Size() == params.hidden, epsilon_, widened == nullptr ? nullptr : widened->MutableData(), squares.get(), sums.get(), buffer.get(), buffer_size, - static_cast(input_width_ == 0 ? params.hidden : input_stride), deterministic_compute); + static_cast(input_width_ == 0 ? params.hidden : input_stride), deterministic_compute, dense_prefill); } return LaunchMixedScaleBranchwiseRMSNorm( Stream(context), x_data, scale == nullptr ? nullptr : scale->DataRaw(), diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu index ebcc0a3dee1ca..c481e05cc8db1 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu @@ -43,7 +43,7 @@ __global__ void BranchwiseSquareKernel(const T* x, float* squares, int64_t count template __global__ void BranchwiseStandardEpilogueKernel( const T* x, const void* scale, int scale_type, const float* sums, T* y, float* widened, - int64_t count, int branches, int hidden, bool shared_scale, float epsilon, int input_row_stride) { + int64_t count, int branches, int hidden, bool shared_scale, float epsilon, int input_row_stride, bool reciprocal_multiply) { const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (index < count) { const int64_t group = index / hidden; @@ -52,7 +52,8 @@ __global__ void BranchwiseStandardEpilogueKernel( const float mean = sums[group] / static_cast(hidden); const float rms = sqrtf(__fadd_rn(mean, epsilon)); const int64_t source_index = Packed ? group * input_row_stride + index % hidden : index; - const float normalized = to_float(x[source_index]) / rms; + const float value = to_float(x[source_index]); + const float normalized = reciprocal_multiply ? __fmul_rn(value, __fdiv_rn(1.0f, rms)) : value / rms; const T rounded = from_float(__fmul_rn(normalized, weight)); y[index] = rounded; if (widened != nullptr) widened[index] = to_float(rounded); @@ -113,7 +114,7 @@ Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const vo int hidden, bool shared_scale, float epsilon, float* widened, float* squares, float* sums, void* reduction_buffer, size_t reduction_buffer_size, int input_row_stride, - bool deterministic_compute) { + bool deterministic_compute, bool reciprocal_multiply) { const int64_t count = static_cast(groups) * hidden; const int64_t blocks = (count + kThreads - 1) / kThreads; ORT_RETURN_IF_NOT(blocks <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); @@ -132,10 +133,10 @@ Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const vo } if (input_row_stride == hidden) { BranchwiseStandardEpilogueKernel<<(blocks), kThreads, 0, stream>>>( - x, scale, scale_type, sums, y, widened, count, branches, hidden, shared_scale, epsilon, input_row_stride); + x, scale, scale_type, sums, y, widened, count, branches, hidden, shared_scale, epsilon, input_row_stride, reciprocal_multiply); } else { BranchwiseStandardEpilogueKernel<<(blocks), kThreads, 0, stream>>>( - x, scale, scale_type, sums, y, widened, count, branches, hidden, shared_scale, epsilon, input_row_stride); + x, scale, scale_type, sums, y, widened, count, branches, hidden, shared_scale, epsilon, input_row_stride, reciprocal_multiply); } return CUDA_CALL(cudaGetLastError()); } @@ -172,7 +173,7 @@ INSTANTIATE(__nv_bfloat16) #define INSTANTIATE_STANDARD(T) \ template Status LaunchStandardBranchwiseRMSNorm(cudaStream_t, const T*, const void*, int, T*, \ - int, int, int, bool, float, float*, float*, float*, void*, size_t, int, bool); + int, int, int, bool, float, float*, float*, float*, void*, size_t, int, bool, bool); INSTANTIATE_STANDARD(float) INSTANTIATE_STANDARD(half) INSTANTIATE_STANDARD(__nv_bfloat16) diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h index 6a80862d62657..c878e8389b6a2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h @@ -20,6 +20,6 @@ Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const vo int hidden, bool shared_scale, float epsilon, float* widened, float* squares, float* sums, void* reduction_buffer, size_t reduction_buffer_size, int input_row_stride, - bool deterministic_compute); + bool deterministic_compute, bool reciprocal_multiply = false); } // namespace onnxruntime::contrib::cuda 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..0ad1fc309c5b7 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc @@ -0,0 +1,115 @@ +// 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" +#include "core/providers/cuda/cuda_type_conversion.h" +#include "core/providers/cuda/shared_inc/fpgeneric.h" +#include "contrib_ops/cuda/bert/branchwise_rms_norm_impl.h" +#include "core/providers/cuda/llm/rotary_embedding_impl.h" +#include "core/providers/cuda/reduction/reduction_functions.h" +#include "contrib_ops/cuda/bert/rotary_embedding_impl.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 && 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), batch = SafeInt(output_batch); + 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_; + if (dim == 256 && groups > 1 && (dense || fp32_rotary_)) { + 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, rotary, cache_rows, scalar, epsilon_, dense, fp32_rotary_); + } + using CudaT = onnxruntime::cuda::OrtToCudaType::type; + auto prepared = GetScratchBuffer(SafeInt(norm_shape.Size()), GetComputeStream(context)); + auto squares = GetScratchBuffer(groups == 1 ? width : 0, GetComputeStream(context)); + auto sums = GetScratchBuffer(groups, GetComputeStream(context)); + const auto buffer_size = onnxruntime::cuda::compute_reduce_matrix_columns_buffer_size(groups, dim); + auto buffer = GetScratchBuffer(buffer_size, GetComputeStream(context)); + ORT_RETURN_IF_ERROR(LaunchStandardBranchwiseRMSNorm( + Stream(context), reinterpret_cast(x), scale->Data(), + ONNX_NAMESPACE::TensorProto::FLOAT, prepared.get(), groups, heads, dim, true, epsilon_, + normalized ? normalized->MutableData() : nullptr, squares.get(), sums.get(), + buffer.get(), buffer_size, stride, context->GetUseDeterministicCompute(), dense)); + const auto launch_rotary = [&]() { + return LaunchRotaryEmbeddingKernel( + Stream(context), reinterpret_cast(output->MutableData()), + reinterpret_cast(prepared.get()), + positions->Data(), nullptr, + reinterpret_cast(cosine->Data()), + reinterpret_cast(sine->Data()), batch, sequence, heads, dim, rotary, + cache_rows, scalar ? 0 : 1, false, GetDeviceProp().maxThreadsPerBlock, false, + !dense && fp32_rotary_); + }; + if (dense) return launch_rotary.template operator()(); + return launch_rotary.template operator()::MappedType>(); +} + +} // 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..42e02f424c1b9 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu @@ -0,0 +1,88 @@ +// 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 rotary_dim, + int cache_rows, bool scalar_position, float epsilon, bool dense_prefill, bool fp32_rotary) { + constexpr int threads = 64; + constexpr int width = 256; + __shared__ float partial[2]; + __shared__ __nv_bfloat16 values[width]; + 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 v[4]; + float sum = 0.f; +#pragma unroll + for (int i = 0; i < 4; ++i) v[i] = __bfloat162float(x[tid + i * threads]); + // Match reduce_square_matrix_columns' four strided values and two-warp shuffle tree. +#pragma unroll + for (int i = 0; i < 4; ++i) sum += v[i] * v[i]; +#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 / float(width), epsilon)); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int column = tid + i * threads; + const float norm = dense_prefill ? __fmul_rn(v[i], __fdiv_rn(1.f, rms)) : v[i] / 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; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int column = tid + i * 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; + } +} + +} +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 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 && input_stride >= 256 && + rotary_dim > 0 && rotary_dim <= 256 && rotary_dim % 2 == 0 && cache_rows > 0 && + (dense_prefill || fp32_rotary), + "RMSNormRotary fast path requires valid dimensions and width 256"); + RMSNormRotaryKernel<<>>( + reinterpret_cast(input), scale, positions, + reinterpret_cast(cosine), reinterpret_cast(sine), + reinterpret_cast<__nv_bfloat16*>(output), normalized, heads, sequence, input_stride, + 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..8836243e0ea45 --- /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 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/bert/rotary_embedding.cc b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc index 722e99b096418..43a486ab388ea 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc @@ -5,6 +5,7 @@ #include "contrib_ops/cpu/bert/rotary_embedding_helper.h" #include "contrib_ops/cuda/bert/rotary_embedding.h" #include "contrib_ops/cuda/bert/rotary_embedding_impl.h" +#include "core/providers/cuda/cuda_type_conversion.h" using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; @@ -24,7 +25,9 @@ namespace cuda { kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("M", DataTypeImpl::GetTensorType()), \ + .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ + .InputMemoryType(OrtMemTypeCPUInput, 4), \ RotaryEmbedding); REGISTER_KERNEL_TYPED(float, 1) @@ -52,6 +55,12 @@ Status RotaryEmbedding::ComputeInternal(OpKernelContext* context) const { const Tensor* position_ids = context->Input(1); const Tensor* cos_cache = context->Input(2); const Tensor* sin_cache = context->Input(3); + const Tensor* numerics = context->Input(4); + ORT_RETURN_IF_NOT(numerics == nullptr || + (numerics->IsDataType() && numerics->Shape().Size() == 1 && + (numerics->Data()[0] == 0 || numerics->Data()[0] == 1)), + "prefill_numerics must be an int32 scalar containing 0 or 1"); + const bool dense_prefill = numerics != nullptr && numerics->Data()[0] == 1; RotaryParameters parameters = {}; ORT_RETURN_IF_ERROR(rotary_embedding_helper::CheckInputs(input, @@ -72,14 +81,16 @@ Status RotaryEmbedding::ComputeInternal(OpKernelContext* context) const { // Launch rotary embedding kernel typedef typename ToCudaType::MappedType CudaT; auto& device_prop = GetDeviceProp(); - return LaunchRotaryEmbeddingKernel( + const auto launch = [&](auto* type_tag) { + using KernelT = std::remove_pointer_t; + return LaunchRotaryEmbeddingKernel( Stream(context), - reinterpret_cast(output->template MutableData()), - reinterpret_cast(input->template Data()), + reinterpret_cast(output->template MutableData()), + reinterpret_cast(input->template Data()), position_ids->Data(), nullptr, // past_sequence_lengths - reinterpret_cast(cos_cache->template Data()), - reinterpret_cast(sin_cache->template Data()), + reinterpret_cast(cos_cache->template Data()), + reinterpret_cast(sin_cache->template Data()), parameters.batch_size, parameters.sequence_length, parameters.num_heads, @@ -90,7 +101,12 @@ Status RotaryEmbedding::ComputeInternal(OpKernelContext* context) const { interleaved, device_prop.maxThreadsPerBlock, parameters.transposed, - use_fp32_arithmetic); + dense_prefill ? false : use_fp32_arithmetic); + }; + if (dense_prefill) { + return launch(static_cast::type*>(nullptr)); + } + return launch(static_cast(nullptr)); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu index 55648cc944249..4415ecc5fef4f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu @@ -246,6 +246,10 @@ template Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, half* out const int max_threads_per_block, const bool is_input_bnsh_format, const bool use_fp32_arithmetic); +template Status LaunchRotaryEmbeddingKernel<__nv_bfloat16>( + cudaStream_t, __nv_bfloat16*, const __nv_bfloat16*, const int64_t*, const int*, + const __nv_bfloat16*, const __nv_bfloat16*, int, int, int, int, int, int, int, bool, int, bool, bool); + template Status LaunchRotaryEmbeddingKernel( cudaStream_t stream, BFloat16* output, const BFloat16* input, const int64_t* position_ids, const int* past_sequence_lengths, const BFloat16* cos_cache, const BFloat16* sin_cache, const int batch_size, const int sequence_length, diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index fe140646872de..9f0e90eb23fc0 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -179,6 +179,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); @@ -549,6 +550,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc index 92af2ee42a95b..cb06ad18f3754 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc @@ -68,6 +68,11 @@ BranchwiseRMSNorm::BranchwiseRMSNorm(const OpKernelInfo& info) } Status BranchwiseRMSNorm::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + const auto* numerics = context.Input(2); + ORT_RETURN_IF_NOT(numerics == nullptr || + (numerics->IsDataType() && numerics->Shape().Size() == 1 && + numerics->Data()[0] == 0), + "BranchwiseRMSNorm: nonzero prefill_numerics requires CUDA"); const auto* x = context.Input(0); const auto* scale = context.Input(1); hyper_connection::StreamShape params; diff --git a/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc b/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc index b4fbdc555a6d5..a3c842dd607fd 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc @@ -17,7 +17,8 @@ ONNX_OPERATOR_KERNEL_EX( kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", WebGpuSupportedFloatTypes()) - .TypeConstraint("M", DataTypeImpl::GetTensorType()), + .TypeConstraint("M", DataTypeImpl::GetTensorType()) + .InputMemoryType(OrtMemTypeCPUInput, 4), RotaryEmbedding); Status RotaryEmbeddingProgram::GenerateShaderCode(ShaderHelper& shader) const { @@ -313,6 +314,11 @@ Status RunRotaryEmbedding(onnxruntime::webgpu::ComputeContext& context, } Status RotaryEmbedding::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + const auto* numerics = context.Input(4); + ORT_RETURN_IF_NOT(numerics == nullptr || + (numerics->IsDataType() && numerics->Shape().Size() == 1 && + numerics->Data()[0] == 0), + "RotaryEmbedding: nonzero prefill_numerics requires CUDA"); const auto* input = context.Input(0); const auto input_shape = input->Shape(); const auto* position_ids = context.Input(1); diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index d233019398d82..4e1cb82d5d49b 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -3517,12 +3517,14 @@ static OpSchema RotaryEmbeddingSchema() { "sin_cache", "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", "T") + .Input(4, "prefill_numerics", "Optional CPU int32 scalar 0/1; one enables dense prefill arithmetic.", "TI", OpSchema::Optional) .Output(0, "output", "tensor with same shape as input.", "T") .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") + .TypeConstraint("TI", {"tensor(int32)"}, "CPU numerical control.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); propagateShapeFromInputToOutput(ctx, 0, 0); @@ -5790,6 +5792,7 @@ static OpSchema BranchwiseRMSNormSchema() { .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", "M", OpSchema::Optional) + .Input(2, "prefill_numerics", "Optional CPU int32 scalar 0/1; one enables dense prefill arithmetic.", "TI", OpSchema::Optional) .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") .Output(1, "Y_fp32", "Optional float32 widening of the already T-rounded Y. " @@ -5802,6 +5805,7 @@ static OpSchema BranchwiseRMSNormSchema() { {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain scale to floating-point tensors.") .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.") + .TypeConstraint("TI", {"tensor(int32)"}, "CPU numerical control.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); const auto start = getAttribute(ctx, "input_start", 0); @@ -5997,5 +6001,88 @@ 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 93b771a425e07..0438896862026 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, 2, BranchwiseRMSNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ScaledSiLU); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPreMix); @@ -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..a4106694adc57 --- /dev/null +++ b/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc @@ -0,0 +1,173 @@ +// 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(); + OpTester original("BranchwiseRMSNorm", 2, kMSDomain); + original.AddAttribute("use_onnx_arithmetic", 1); + original.AddAttribute("input_width", width); + original.AddAttribute("input_start", start); + original.AddAttribute("epsilon", 1e-6f); + original.AddInput("X", input_shape, input); + original.AddInput("scale", {width}, scale); + original.AddInput("mode", {1}, {dense}); + original.AddOutput("Y", norm_shape, norm); + original.AddOutput("Y_fp32", norm_shape, normalized); + original.SetCustomOutputVerifier([&](const std::vector& values, const std::string&) { + if (norm.empty()) return; + const auto& y = values[0].Get(); + norm.assign(y.Data(), y.Data() + norm.size()); + const auto& widened = values[1].Get(); + normalized.assign(widened.Data(), widened.Data() + normalized.size()); + }); + run(original); + OpTester rope("RotaryEmbedding", 2, kMSDomain); + rope.AddAttribute("num_heads", heads); + rope.AddAttribute("rotary_embedding_dim", rotary); + rope.AddAttribute("use_fp32_arithmetic", fp32 ? 1 : 0); + rope.AddInput("X", output_shape, norm); + rope.AddInput("position", position_shape, positions); + rope.AddInput("cos", {cache_rows, rotary / 2}, cosine); + rope.AddInput("sin", {cache_rows, rotary / 2}, sine); + rope.AddInput("mode", {1}, {dense}); + rope.AddOutput("Y", output_shape, expected); + rope.SetCustomOutputVerifier([&](const std::vector& values, const std::string&) { + if (expected.empty()) return; + const auto& y = values[0].Get(); + expected.assign(y.Data(), y.Data() + expected.size()); + }); + if (sequence > 0) run(rope); + 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 From d68e93b5f3c4828579bd46360ba3430dc5427f92 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Oct 2026 23:14:10 +0000 Subject: [PATCH 5/8] [CUDA] Remove unused RotaryEmbedding(ver 2) from PR #33469 and update RMSNormRotary/BranchwiseRMSNorm docs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 78 +++++++++- docs/OperatorKernels.md | 10 +- .../contrib_ops/cpu/bert/rotary_embedding.cc | 5 - .../contrib_ops/cuda/bert/rotary_embedding.cc | 47 ++---- .../contrib_ops/cuda/bert/rotary_embedding.h | 1 - .../contrib_ops/cuda/cuda_contrib_kernels.cc | 6 - .../webgpu/bert/rotary_embedding.cc | 12 +- .../core/graph/contrib_ops/bert_defs.cc | 114 ++++++-------- onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 - .../optimizer/group_query_attention_fusion.cc | 3 - .../contrib_ops/rotary_embedding_op_test.cc | 145 ------------------ 11 files changed, 143 insertions(+), 280 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 0ad97f90ebea8..96d87b1130dcd 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -111,6 +111,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 @@ -925,13 +926,15 @@ Other versions of this operator: com.microsoft.BranchwiseRMSNorm-1
Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. Supported values are 0 (default) and 1.
-#### Inputs (1 - 2) +#### Inputs (1 - 3)
X : T
Grouped (..., C, H) or flattened (..., C * H) input.
scale (optional) : M
Optional scale with shape (C * H), (C, H), or (H).
+
prefill_numerics (optional) : TI
+
Optional CPU int32 scalar 0/1; one enables dense prefill arithmetic.
#### Outputs (1 - 2) @@ -952,6 +955,8 @@ Other versions of this operator: com.microsoft.BranchwiseRMSNorm-1
Constrain scale to floating-point tensors.
F : tensor(float)
Widened output is float32.
+
TI : tensor(int32)
+
CPU numerical control.
@@ -6942,6 +6947,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 (default is 1e-06)
+
Finite nonnegative RMS epsilon.
+
flatten_batch : int (default is 0)
+
One flattens B*S to sequence and outputs [1,B*S,heads*width].
+
input_start : int (default is 0)
+
Nonnegative start of normalized last-axis slice.
+
input_width : int (default is 0)
+
Normalized width; zero uses the full last dimension.
+
rotary_embedding_dim : int
+
Positive even rotary width, no larger than normalized width.
+
use_fp32_arithmetic : int (default is 1)
+
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` @@ -7197,9 +7267,7 @@ This version of the operator has been available since version 1 of the 'com.micr #### Version -This version of the operator has been available since version 2 of the 'com.microsoft' operator set. - -Other versions of this operator: com.microsoft.RotaryEmbedding-1 +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. #### Attributes @@ -7214,8 +7282,6 @@ Other versions of this operator: com.microsoft.RotaryEmbedding-1
Rotary embedding dimension. Default value is 0.
scale : float
Custom scale will be used if specified. Default value is 1.0
-
use_fp32_arithmetic : int (default is 0)
-
Widen input and cache values to float32, round each multiplication and addition in float32, then convert the result to the input type. Values are 0 (legacy arithmetic) or 1. Currently supported by CUDA only.
#### Inputs diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 46971132fefeb..01d2daf6a684b 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -575,7 +575,7 @@ The **OpSet Version** column uses the following notation: |BeamSearch|*in* input_ids:**F**
*in* max_length:**I**
*in* min_length:**I**
*in* num_beams:**I**
*in* num_return_sequences:**I**
*in* length_penalty:**T**
*in* repetition_penalty:**T**
*in* vocab_mask:**M**
*in* prefix_vocab_mask:**M**
*in* attention_mask:**I**
*in* decoder_input_ids:**I**
*in* logits_processor:**I**
*out* sequences:**I**
*out* sequences_scores:**T**
*out* scores:**T**|1+|**T** = tensor(float)| |BiasGelu|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float)| |BifurcationDetector|*in* src_tokens:**T**
*in* cur_tokens:**T**
*in* prev_suffix_match_idx:**T**
*in* pred_tokens:**T**
*out* tokens:**T**
*out* suffix_match_idx:**T**|1+|**T** = tensor(int64)| -|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**
*out* Y_fp32:**F**|1+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*in* prefill_numerics:**TI**
*out* Y:**T**
*out* Y_fp32:**F**|1+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |CDist|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(double), tensor(float)| |CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**|1+|**T** = tensor(float)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| @@ -1088,8 +1088,8 @@ The **OpSet Version** column uses the following notation: |BiasSplitGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |BitmaskBiasDropout|*in* data:**T**
*in* bias:**T**
*in* residual:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| |BitmaskDropout|*in* data:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| -|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**
*out* Y_fp32:**F**|2+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| -|||1+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*in* prefill_numerics:**TI**
*out* Y:**T**
*out* Y_fp32:**F**|2+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)| +|||1+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)| |CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |ComplexMul|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ComplexMulConj|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| @@ -1150,12 +1150,12 @@ 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)| |Rfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| -|RotaryEmbedding|*in* input:**T**
*in* position_ids:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**|2+|**M** = tensor(int64)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| -|||1+|**M** = tensor(int64)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|RotaryEmbedding|*in* input:**T**
*in* position_ids:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**|1+|**M** = tensor(int64)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |Sampling|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* repetition_penalty:**T**
*in* vocab_mask:**I**
*in* prefix_vocab_mask:**I**
*in* attention_mask:**I**
*in* presence_mask:**I**
*in* seed:**I**
*out* sequences:**I**
*out* filtered_logits:**T**|1+|**T** = tensor(float), tensor(float16)| |ScaledSiLU|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |SkipGroupNorm|*in* X:**T**
*in* gamma:**M**
*in* beta:**M**
*in* skip:**T**
*in* bias:**T**
*out* Y:**T**
*out* S:**T**|1+|**T** = tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc b/onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc index c9b9e5a8626f7..682c254840e91 100644 --- a/onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc +++ b/onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc @@ -184,11 +184,6 @@ template Status RunRotaryEmbedding(concurrency::ThreadPool* tp, Rotar template Status RotaryEmbedding::Compute(OpKernelContext* context) const { - const auto* numerics = context->Input(4); - ORT_RETURN_IF_NOT(numerics == nullptr || - (numerics->IsDataType() && numerics->Shape().Size() == 1 && - numerics->Data()[0] == 0), - "RotaryEmbedding: nonzero prefill_numerics requires CUDA"); const Tensor* input = context->Input(0); const Tensor* position_ids = context->Input(1); const Tensor* cos_cache = context->Input(2); diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc index 43a486ab388ea..0b38c6f0e5484 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc @@ -5,7 +5,6 @@ #include "contrib_ops/cpu/bert/rotary_embedding_helper.h" #include "contrib_ops/cuda/bert/rotary_embedding.h" #include "contrib_ops/cuda/bert/rotary_embedding_impl.h" -#include "core/providers/cuda/cuda_type_conversion.h" using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; @@ -16,26 +15,21 @@ namespace onnxruntime { namespace contrib { namespace cuda { -#define REGISTER_KERNEL_TYPED(T, Version) \ +#define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ RotaryEmbedding, \ kMSDomain, \ - Version, \ + 1, \ T, \ kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ - .InputMemoryType(OrtMemTypeCPUInput, 4), \ + .TypeConstraint("M", DataTypeImpl::GetTensorType()), \ RotaryEmbedding); -REGISTER_KERNEL_TYPED(float, 1) -REGISTER_KERNEL_TYPED(MLFloat16, 1) -REGISTER_KERNEL_TYPED(BFloat16, 1) -REGISTER_KERNEL_TYPED(float, 2) -REGISTER_KERNEL_TYPED(MLFloat16, 2) -REGISTER_KERNEL_TYPED(BFloat16, 2) +REGISTER_KERNEL_TYPED(float) +REGISTER_KERNEL_TYPED(MLFloat16) +REGISTER_KERNEL_TYPED(BFloat16) template RotaryEmbedding::RotaryEmbedding(const OpKernelInfo& info) : CudaKernel(info) { @@ -44,9 +38,6 @@ RotaryEmbedding::RotaryEmbedding(const OpKernelInfo& info) : CudaKernel(info) num_heads = static_cast(info.GetAttrOrDefault("num_heads", 0)); interleaved = (info.GetAttrOrDefault("interleaved", 0) == 1); is_packed_batching = (info.GetAttrOrDefault("is_packed_batching", 0) == 1); - const auto mode = info.GetAttrOrDefault("use_fp32_arithmetic", 0); - ORT_ENFORCE(mode == 0 || mode == 1, "use_fp32_arithmetic must be 0 or 1"); - use_fp32_arithmetic = mode == 1; } template @@ -55,12 +46,6 @@ Status RotaryEmbedding::ComputeInternal(OpKernelContext* context) const { const Tensor* position_ids = context->Input(1); const Tensor* cos_cache = context->Input(2); const Tensor* sin_cache = context->Input(3); - const Tensor* numerics = context->Input(4); - ORT_RETURN_IF_NOT(numerics == nullptr || - (numerics->IsDataType() && numerics->Shape().Size() == 1 && - (numerics->Data()[0] == 0 || numerics->Data()[0] == 1)), - "prefill_numerics must be an int32 scalar containing 0 or 1"); - const bool dense_prefill = numerics != nullptr && numerics->Data()[0] == 1; RotaryParameters parameters = {}; ORT_RETURN_IF_ERROR(rotary_embedding_helper::CheckInputs(input, @@ -81,16 +66,14 @@ Status RotaryEmbedding::ComputeInternal(OpKernelContext* context) const { // Launch rotary embedding kernel typedef typename ToCudaType::MappedType CudaT; auto& device_prop = GetDeviceProp(); - const auto launch = [&](auto* type_tag) { - using KernelT = std::remove_pointer_t; - return LaunchRotaryEmbeddingKernel( + return LaunchRotaryEmbeddingKernel( Stream(context), - reinterpret_cast(output->template MutableData()), - reinterpret_cast(input->template Data()), + reinterpret_cast(output->template MutableData()), + reinterpret_cast(input->template Data()), position_ids->Data(), nullptr, // past_sequence_lengths - reinterpret_cast(cos_cache->template Data()), - reinterpret_cast(sin_cache->template Data()), + reinterpret_cast(cos_cache->template Data()), + reinterpret_cast(sin_cache->template Data()), parameters.batch_size, parameters.sequence_length, parameters.num_heads, @@ -100,13 +83,7 @@ Status RotaryEmbedding::ComputeInternal(OpKernelContext* context) const { parameters.position_ids_format, interleaved, device_prop.maxThreadsPerBlock, - parameters.transposed, - dense_prefill ? false : use_fp32_arithmetic); - }; - if (dense_prefill) { - return launch(static_cast::type*>(nullptr)); - } - return launch(static_cast(nullptr)); + parameters.transposed); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h index 4487746579a9a..55f654a2cea8f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h @@ -23,7 +23,6 @@ class RotaryEmbedding final : public CudaKernel { int rotary_embedding_dim; bool interleaved; bool is_packed_batching; - bool use_fp32_arithmetic; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 9f0e90eb23fc0..9ef025f366433 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -160,9 +160,6 @@ class CUDA_ONNX_OP_TYPED_CLASS_NAME(1, MLFloat16, ParametricSoftplus); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, RotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, RotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, RotaryEmbedding); -class CUDA_MS_OP_TYPED_CLASS_NAME(2, float, RotaryEmbedding); -class CUDA_MS_OP_TYPED_CLASS_NAME(2, MLFloat16, RotaryEmbedding); -class CUDA_MS_OP_TYPED_CLASS_NAME(2, BFloat16, RotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MRotaryEmbedding); @@ -531,9 +528,6 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc b/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc index a3c842dd607fd..f324cfa2e04cf 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc @@ -17,8 +17,7 @@ ONNX_OPERATOR_KERNEL_EX( kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", WebGpuSupportedFloatTypes()) - .TypeConstraint("M", DataTypeImpl::GetTensorType()) - .InputMemoryType(OrtMemTypeCPUInput, 4), + .TypeConstraint("M", DataTypeImpl::GetTensorType()), RotaryEmbedding); Status RotaryEmbeddingProgram::GenerateShaderCode(ShaderHelper& shader) const { @@ -158,7 +157,7 @@ Status FusedQKRotaryEmbeddingProgram::GenerateShaderCode(ShaderHelper& shader) c << " let q_v = f32(" << q_input.GetByOffset("q_head_base + c") << ");\n" << " q_sumsq = q_sumsq + q_v * q_v;\n" << " }\n" - << " let q_inv_rms = q_input_element_t(inverseSqrt(q_sumsq / f32(uniforms.head_size) + uniforms.qk_norm_epsilon));\n" + << " let q_inv_rms = q_input_element_t(inverseSqrt(q_sumsq / f32(uniforms.head_size) + uniforms.q_norm_epsilon));\n" << " let k_head_base = bsnh[0] * uniforms.k_input_output_stride[0]\n" << " + bsnh[1] * uniforms.k_input_output_stride[1]\n" << " + bsnh[2] * uniforms.k_input_output_stride[2];\n" @@ -169,7 +168,7 @@ Status FusedQKRotaryEmbeddingProgram::GenerateShaderCode(ShaderHelper& shader) c << " let k_v = f32(" << k_input.GetByOffset("k_head_base + c") << ");\n" << " k_sumsq = k_sumsq + k_v * k_v;\n" << " }\n" - << " k_inv_rms = k_input_element_t(inverseSqrt(k_sumsq / f32(uniforms.head_size) + uniforms.qk_norm_epsilon));\n" + << " k_inv_rms = k_input_element_t(inverseSqrt(k_sumsq / f32(uniforms.head_size) + uniforms.k_norm_epsilon));\n" << " }\n"; } @@ -314,11 +313,6 @@ Status RunRotaryEmbedding(onnxruntime::webgpu::ComputeContext& context, } Status RotaryEmbedding::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { - const auto* numerics = context.Input(4); - ORT_RETURN_IF_NOT(numerics == nullptr || - (numerics->IsDataType() && numerics->Shape().Size() == 1 && - numerics->Data()[0] == 0), - "RotaryEmbedding: nonzero prefill_numerics requires CUDA"); const auto* input = context.Input(0); const auto input_shape = input->Shape(); const auto* position_ids = context.Input(1); diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 4e1cb82d5d49b..8025a4969226b 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -3476,70 +3476,58 @@ constexpr const char* RotaryEmbedding_ver1_doc = R"DOC( RotaryEmbedding is the implementation of rotary positional embeddings (RoPE). The positions are represented as rotation matrices that are multiplied to query and key before the inner product of query and key is taken. )DOC"; -static OpSchema RotaryEmbeddingSchema() { - return OpSchema() - .SetDoc(RotaryEmbedding_ver1_doc) - .Attr("scale", - "Custom scale will be used if specified. Default value is 1.0", - AttributeProto::FLOAT, - OPTIONAL_VALUE) - .Attr("interleaved", - "Indicates whether the input has real and imaginary parts interleaved. " - "Default value is 0 (False), meaning the first half of the input consists of real values " - "and the second half consists of imaginary values.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("rotary_embedding_dim", - "Rotary embedding dimension. Default value is 0.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("num_heads", - "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("is_packed_batching", - "ragged batch inputs or not. Default value is 0", - AttributeProto::INT, - OPTIONAL_VALUE) - .Input(0, - "input", - "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", - "T") - .Input(1, - "position_ids", - "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", - "M") - .Input(2, - "cos_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Input(3, - "sin_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Input(4, "prefill_numerics", "Optional CPU int32 scalar 0/1; one enables dense prefill arithmetic.", "TI", OpSchema::Optional) - .Output(0, - "output", - "tensor with same shape as input.", - "T") - .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") - .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") - .TypeConstraint("TI", {"tensor(int32)"}, "CPU numerical control.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - propagateShapeFromInputToOutput(ctx, 0, 0); - }); -} - -ONNX_MS_OPERATOR_SET_SCHEMA(RotaryEmbedding, 1, RotaryEmbeddingSchema()); ONNX_MS_OPERATOR_SET_SCHEMA( - RotaryEmbedding, 2, - RotaryEmbeddingSchema() - .Attr("use_fp32_arithmetic", - "Widen input and cache values to float32, round each multiplication and addition in float32, " - "then convert the result to the input type. Values are 0 (legacy arithmetic) or 1. " - "Currently supported by CUDA only.", - AttributeProto::INT, static_cast(0))); + RotaryEmbedding, 1, + OpSchema() + .SetDoc(RotaryEmbedding_ver1_doc) + .Attr("scale", + "Custom scale will be used if specified. Default value is 1.0", + AttributeProto::FLOAT, + OPTIONAL_VALUE) + .Attr("interleaved", + "Indicates whether the input has real and imaginary parts interleaved. " + "Default value is 0 (False), meaning the first half of the input consists of real values " + "and the second half consists of imaginary values.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("rotary_embedding_dim", + "Rotary embedding dimension. Default value is 0.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("num_heads", + "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("is_packed_batching", + "ragged batch inputs or not. Default value is 0", + AttributeProto::INT, + OPTIONAL_VALUE) + .Input(0, + "input", + "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", + "T") + .Input(1, + "position_ids", + "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", + "M") + .Input(2, + "cos_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Input(3, + "sin_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Output(0, + "output", + "tensor with same shape as input.", + "T") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") + .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateShapeFromInputToOutput(ctx, 0, 0); + })); constexpr const char* MRotaryEmbedding_ver1_doc = R"DOC( MRotaryEmbedding is the fused implementation of Multimodal Rotary Positional Embeddings (M-RoPE) used by the diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 0438896862026..9d4bb285bc560 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -116,7 +116,6 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RemovePadding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RestorePadding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Rfft); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RotaryEmbedding); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 2, RotaryEmbedding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MRotaryEmbedding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GemmaRotaryEmbedding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SampleOp); @@ -255,7 +254,6 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); - fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/core/optimizer/group_query_attention_fusion.cc b/onnxruntime/core/optimizer/group_query_attention_fusion.cc index d1903c297f210..6506dfae7898d 100644 --- a/onnxruntime/core/optimizer/group_query_attention_fusion.cc +++ b/onnxruntime/core/optimizer/group_query_attention_fusion.cc @@ -352,9 +352,6 @@ static bool TryGetRotaryEmbeddingArgs(Node& rotary_node, RotaryEmbeddingArgs& ar if (rotary_node.Domain() != kMSDomain) { return false; } - if (GetIntAttributeOrDefault(rotary_node, "use_fp32_arithmetic", 0) != 0) { - return false; - } args.interleaved = GetIntAttributeOrDefault(rotary_node, "interleaved", 0); args.rotary_embedding_dim = GetIntAttributeOrDefault(rotary_node, "rotary_embedding_dim", 0); diff --git a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc index a6957e18bc194..658c8695643e6 100644 --- a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc +++ b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc @@ -2,8 +2,6 @@ // Licensed under the MIT License. #include -#include -#include #include #include "gtest/gtest.h" #include "contrib_ops/cpu/bert/rotary_embedding.h" @@ -23,149 +21,6 @@ enum class TensorType { }; } // anonymous namespace -class FloatRotaryChainTester : public OpTester { - public: - FloatRotaryChainTester(int type, int heads, int rotary_dim, bool interleaved) - : OpTester("RotaryEmbedding", 1, kMSDomain, false), - type_(type), - heads_(heads), - rotary_dim_(rotary_dim), - interleaved_(interleaved) {} - - void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, - std::vector>&) override { - auto value = [&](const std::string& name) { return &graph.GetOrCreateNodeArg(name, nullptr); }; - std::vector widened = inputs; - for (size_t index : {size_t{0}, size_t{2}, size_t{3}}) { - const auto name = "float_" + std::to_string(index); - widened[index] = value(name); - graph.AddNode(name, "Cast", "", {inputs[index]}, {widened[index]}) - .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); - } - auto* result = value("float_result"); - auto& rotary = graph.AddNode("rotary", "RotaryEmbedding", "", widened, {result}, nullptr, kMSDomain); - rotary.AddAttribute("num_heads", static_cast(heads_)); - rotary.AddAttribute("rotary_embedding_dim", static_cast(rotary_dim_)); - rotary.AddAttribute("interleaved", static_cast(interleaved_)); - graph.AddNode("rounded", "Cast", "", {result}, {outputs[0]}) - .AddAttribute("to", static_cast(type_)); - graph.SetOutputs(outputs); - } - - private: - void CreateModelToTest(const ModelOptions& options, Model*& model) override { - model = &BuildModel({{kOnnxDomain, 17}}, options); - ORT_THROW_IF_ERROR(model->MainGraph().Resolve()); - } - - int type_, heads_, rotary_dim_; - bool interleaved_; -}; - -template -void TestFloatRotaryArithmetic(int type) { - if (!HasCudaEnvironment(800)) { - GTEST_SKIP() << "CUDA SM80 is required"; - } - for (int sequence : {1, 2, 9}) { - for (int width : {64, 256}) { - for (bool transposed : {false, true}) { - for (bool interleaved : {false, true}) { - SCOPED_TRACE("sequence=" + std::to_string(sequence) + " width=" + std::to_string(width) + - " transposed=" + std::to_string(transposed) + " interleaved=" + std::to_string(interleaved)); - constexpr int batch = 2, heads = 3; - const int rotary_dim = 64; - const int cache_rows = batch * sequence + 3; - std::vector shape = transposed ? std::vector{batch, heads, sequence, width} - : std::vector{batch, sequence, heads * width}; - std::vector input(batch * sequence * heads * width), cosine(cache_rows * rotary_dim / 2), - sine(cosine.size()); - const bool scalar_position = !transposed; - std::vector positions(scalar_position ? 1 : batch * sequence); - for (size_t i = 0; i < input.size(); ++i) input[i] = T(std::sin(i * 0.71f) * (1 + i % 17)); - for (size_t i = 0; i < cosine.size(); ++i) { - cosine[i] = T(std::cos(i * 0.13f)); - sine[i] = T(std::sin(i * 0.13f)); - } - for (size_t i = 0; i < positions.size(); ++i) positions[i] = scalar_position ? 1 : (i * 7) % cache_rows; - FloatRotaryChainTester reference(type, heads, rotary_dim, interleaved); - OpTester fused("RotaryEmbedding", 2, kMSDomain, false); - fused.AddAttribute("num_heads", heads); - fused.AddAttribute("rotary_embedding_dim", rotary_dim); - fused.AddAttribute("interleaved", interleaved); - fused.AddAttribute("use_fp32_arithmetic", 1); - std::vector expected; - for (OpTester* tester : {static_cast(&reference), &fused}) { - tester->AddInput("input", shape, input); - tester->AddInput("positions", scalar_position ? std::vector{1} : std::vector{batch, sequence}, positions); - tester->AddInput("cosine", {cache_rows, rotary_dim / 2}, cosine); - tester->AddInput("sine", {cache_rows, rotary_dim / 2}, sine); - tester->AddOutput("output", shape, std::vector(input.size())); - SessionOptions options; - options.graph_optimization_level = TransformerLevel::Default; - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - tester->Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); - auto fetches = tester->GetFetches(); - ASSERT_EQ(fetches.size(), 1); - if (expected.empty()) { - expected = std::move(fetches); - } else { - const auto& before = expected[0].Get(); - const auto& after = fetches[0].Get(); - ASSERT_EQ(before.SizeInBytes(), after.SizeInBytes()); - ASSERT_EQ(std::memcmp(before.DataRaw(), after.DataRaw(), before.SizeInBytes()), 0); - } - } - } - } - } - } -} - -TEST(RotaryEmbeddingTest, FloatArithmeticMatchesWidenRotateRoundFloat16) { - TestFloatRotaryArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); -} - -TEST(RotaryEmbeddingTest, FloatArithmeticMatchesWidenRotateRoundBFloat16) { - TestFloatRotaryArithmetic(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); -} - -TEST(RotaryEmbeddingTest, Version2LegacyDefaultAndInvalidMode) { - if (!HasCudaEnvironment(800)) { - GTEST_SKIP() << "CUDA SM80 is required"; - } - for (int mode : {0, 2}) { - OpTester test("RotaryEmbedding", 2, kMSDomain); - if (mode != 0) test.AddAttribute("use_fp32_arithmetic", mode); - test.AddAttribute("num_heads", 1); - test.AddInput("input", {1, 1, 4}, ToFloat16({1, 2, 3, 4})); - test.AddInput("positions", {1}, {0}); - test.AddInput("cos", {1, 2}, ToFloat16({1, 1})); - test.AddInput("sin", {1, 2}, ToFloat16({0, 0})); - test.AddOutput("output", {1, 1, 4}, ToFloat16({1, 2, 3, 4})); - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, - mode == 0 ? "" : "use_fp32_arithmetic must be 0 or 1", {}, nullptr, &providers); - } -} - -TEST(RotaryEmbeddingTest, Version2DoesNotFallBackToLegacyCpuArithmetic) { - OpTester test("RotaryEmbedding", 2, kMSDomain); - test.AddAttribute("use_fp32_arithmetic", 1); - test.AddAttribute("num_heads", 1); - test.AddInput("input", {1, 1, 4}, {1, 2, 3, 4}); - test.AddInput("positions", {1}, {0}); - test.AddInput("cos", {1, 2}, {1, 1}); - test.AddInput("sin", {1, 2}, {0, 0}); - test.AddOutput("output", {1, 1, 4}, {1, 2, 3, 4}); - std::vector> providers; - providers.push_back(DefaultCpuExecutionProvider()); - test.Run(OpTester::ExpectResult::kExpectFailure, "Could not find an implementation for RotaryEmbedding(2)", - {}, nullptr, &providers); -} - static void RunTest( const std::vector& input_data, const std::vector& position_ids, From 52b3e0d7f7582fd569bf736832d0a01a8ef80072 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Oct 2026 23:41:03 +0000 Subject: [PATCH 6/8] Apply lintrunner formatting Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/branchwise_rms_norm.cc | 2 +- .../contrib_ops/cuda/bert/rms_norm_rotary.cc | 14 +++++++------- .../contrib_ops/cuda/bert/rms_norm_rotary.cu | 10 +++++----- .../contrib_ops/cuda/bert/rms_norm_rotary.h | 8 ++++---- onnxruntime/core/graph/contrib_ops/bert_defs.cc | 8 +++++--- .../test/contrib_ops/rms_norm_rotary_test.cc | 13 ++++++++----- 6 files changed, 30 insertions(+), 25 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc index bc6b356aba41d..47aa148359aa1 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc @@ -22,7 +22,7 @@ using namespace onnxruntime::cuda; .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ .TypeConstraint("F", DataTypeImpl::GetTensorType()) \ .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ - .InputMemoryType(OrtMemTypeCPUInput, 2) \ + .InputMemoryType(OrtMemTypeCPUInput, 2) \ .TypeConstraint("M", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}), \ diff --git a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc index 0ad1fc309c5b7..369489f585baa 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc @@ -100,13 +100,13 @@ Status RMSNormRotary::ComputeInternal(OpKernelContext* context) const { buffer.get(), buffer_size, stride, context->GetUseDeterministicCompute(), dense)); const auto launch_rotary = [&]() { return LaunchRotaryEmbeddingKernel( - Stream(context), reinterpret_cast(output->MutableData()), - reinterpret_cast(prepared.get()), - positions->Data(), nullptr, - reinterpret_cast(cosine->Data()), - reinterpret_cast(sine->Data()), batch, sequence, heads, dim, rotary, - cache_rows, scalar ? 0 : 1, false, GetDeviceProp().maxThreadsPerBlock, false, - !dense && fp32_rotary_); + Stream(context), reinterpret_cast(output->MutableData()), + reinterpret_cast(prepared.get()), + positions->Data(), nullptr, + reinterpret_cast(cosine->Data()), + reinterpret_cast(sine->Data()), batch, sequence, heads, dim, rotary, + cache_rows, scalar ? 0 : 1, false, GetDeviceProp().maxThreadsPerBlock, false, + !dense && fp32_rotary_); }; if (dense) return launch_rotary.template operator()(); return launch_rotary.template operator()::MappedType>(); diff --git a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu index 42e02f424c1b9..42acb0ff84197 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu @@ -67,12 +67,12 @@ __global__ void RMSNormRotaryKernel( } } -} +} // 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 rotary_dim, int cache_rows, bool scalar_position, - float epsilon, bool dense_prefill, bool fp32_rotary) { + 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 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 && input_stride >= 256 && rotary_dim > 0 && rotary_dim <= 256 && rotary_dim % 2 == 0 && cache_rows > 0 && (dense_prefill || fp32_rotary), diff --git a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h index 8836243e0ea45..f6220658b8413 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h @@ -5,10 +5,10 @@ #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 rotary_dim, int cache_rows, bool scalar_position, - float epsilon, bool dense_prefill, bool fp32_rotary); + 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 rotary_dim, int cache_rows, bool scalar_position, + float epsilon, bool dense_prefill, bool fp32_rotary); class RMSNormRotary final : public onnxruntime::cuda::CudaKernel { public: diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 49a8d278c45d2..028bd81694fcf 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -6216,8 +6216,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .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) + .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.") @@ -6250,7 +6252,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( 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())) { + scale.dim(0).dim_value() != normalized.dim(3).dim_value())) { fail_shape_inference("RMSNormRotary scale must have shape [width]."); } } diff --git a/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc b/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc index a4106694adc57..73ce50f748149 100644 --- a/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc +++ b/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc @@ -77,8 +77,7 @@ TEST(RMSNormRotaryTest, RMSNormRotaryExactUnfusedParity) { 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 position_shape = scalar ? std::vector{1} : std::vector{output_batch, output_sequence}; const auto run = [](OpTester& test) { auto providers = CudaProvider(); ASSERT_FALSE(providers.empty()); @@ -162,9 +161,13 @@ TEST(RMSNormRotaryTest, RMSNormRotaryRejectsInvalidContracts) { 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", + 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); } } From 12caf709326d9adb443675846a23868852873aac Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 00:30:28 +0000 Subject: [PATCH 7/8] [CUDA] Revert unused BranchwiseRMSNorm(ver 2) and unfused RMSNormRotary fallback from PR #33469 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 35 +- docs/OperatorKernels.md | 5 +- .../cpu/bert/branchwise_rms_norm.cc | 27 +- .../cpu/bert/branchwise_rms_norm.h | 1 - .../cuda/bert/branchwise_rms_norm.cc | 80 +--- .../cuda/bert/branchwise_rms_norm.h | 3 - .../cuda/bert/branchwise_rms_norm_impl.cu | 72 ---- .../cuda/bert/branchwise_rms_norm_impl.h | 8 - .../contrib_ops/cuda/bert/rms_norm_rotary.cc | 50 +-- .../contrib_ops/cuda/bert/rms_norm_rotary.cu | 40 +- .../contrib_ops/cuda/bert/rms_norm_rotary.h | 2 +- .../cuda/bert/rotary_embedding_impl.cu | 66 ++-- .../cuda/bert/rotary_embedding_impl.h | 6 +- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 6 - .../webgpu/bert/branchwise_rms_norm.cc | 11 +- .../core/graph/contrib_ops/bert_defs.cc | 105 +---- onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 - .../cuda/reduction/reduction_functions.cu | 31 +- .../cuda/reduction/reduction_functions.h | 5 - onnxruntime/core/session/environment.cc | 2 +- .../contrib_ops/hyper_connection_op_test.cc | 368 ------------------ .../test/contrib_ops/rms_norm_rotary_test.cc | 91 +++-- .../graph_transform_test_layernorm.cc | 30 +- .../core/graph/training_op_defs.cc | 2 +- 24 files changed, 161 insertions(+), 887 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 1c5043a8947a7..db014df0124b8 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -899,55 +899,34 @@ This version of the operator has been available since version 1 of the 'com.micr (..., C, H), or flattened shape (..., C * H) when num_branches is specified. The optional scale may have shape (C * H), (C, H), or (H). Arithmetic is performed in float32 and the result is converted to T. - With use_onnx_arithmetic=1, compute FP32 square, mean, epsilon addition, square - root, division, and scale multiplication with separate rounding boundaries. - CUDA reuses the standard fast last-axis ReduceMean summation order by default. - Deterministic CUDA sessions instead use a fixed-order reduction, which may - round differently from the default reduction or deterministic ReduceMean. - This mode does not substitute reciprocal square root. - Optional Y_fp32 losslessly widens the already T-rounded Y to float32; it does - not bypass the output rounding boundary and requires use_onnx_arithmetic=1. - WebGPU currently supports only the default mode without Y_fp32. #### Version -This version of the operator has been available since version 2 of the 'com.microsoft' operator set. - -Other versions of this operator: com.microsoft.BranchwiseRMSNorm-1 +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. #### Attributes
-
epsilon : float (default is 1e-05)
+
epsilon : float
Epsilon added before reciprocal square root.
-
input_start : int (default is 0)
-
CUDA-only starting offset in the last input dimension.
-
input_width : int (default is 0)
-
CUDA-only width selected from the last input dimension before normalization. Zero retains the full input. Requires grouped input and use_onnx_arithmetic=1. Outputs retain the leading dimensions and use this last dimension.
-
num_branches : int (default is 0)
+
num_branches : int
Number of branches for flattened input. Omit or set to zero for grouped input.
-
use_onnx_arithmetic : int (default is 0)
-
Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. Supported values are 0 (default) and 1.
-#### Inputs (1 - 3) +#### Inputs (1 - 2)
X : T
Grouped (..., C, H) or flattened (..., C * H) input.
scale (optional) : M
Optional scale with shape (C * H), (C, H), or (H).
-
prefill_numerics (optional) : TI
-
Optional CPU int32 scalar 0/1; one enables dense prefill arithmetic.
-#### Outputs (1 - 2) +#### Outputs
Y : T
RMS-normalized output with the same shape as X.
-
Y_fp32 (optional) : F
-
Optional float32 widening of the already T-rounded Y. Requires use_onnx_arithmetic=1.
#### Type Constraints @@ -957,10 +936,6 @@ Other versions of this operator: com.microsoft.BranchwiseRMSNorm-1
Constrain input and output to floating-point tensors.
M : tensor(float), tensor(float16), tensor(bfloat16)
Constrain scale to floating-point tensors.
-
F : tensor(float)
-
Widened output is float32.
-
TI : tensor(int32)
-
CPU numerical control.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index d5a1e16eaffe3..6d084f06a97dc 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -583,7 +583,7 @@ The **OpSet Version** column uses the following notation: |BeamSearch|*in* input_ids:**F**
*in* max_length:**I**
*in* min_length:**I**
*in* num_beams:**I**
*in* num_return_sequences:**I**
*in* length_penalty:**T**
*in* repetition_penalty:**T**
*in* vocab_mask:**M**
*in* prefix_vocab_mask:**M**
*in* attention_mask:**I**
*in* decoder_input_ids:**I**
*in* logits_processor:**I**
*out* sequences:**I**
*out* sequences_scores:**T**
*out* scores:**T**|1+|**T** = tensor(float)| |BiasGelu|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float)| |BifurcationDetector|*in* src_tokens:**T**
*in* cur_tokens:**T**
*in* prev_suffix_match_idx:**T**
*in* pred_tokens:**T**
*out* tokens:**T**
*out* suffix_match_idx:**T**|1+|**T** = tensor(int64)| -|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*in* prefill_numerics:**TI**
*out* Y:**T**
*out* Y_fp32:**F**|1+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |CDist|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(double), tensor(float)| |CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**|1+|**T** = tensor(float)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| @@ -1102,8 +1102,7 @@ The **OpSet Version** column uses the following notation: |BiasSplitGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |BitmaskBiasDropout|*in* data:**T**
*in* bias:**T**
*in* residual:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| |BitmaskDropout|*in* data:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| -|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*in* prefill_numerics:**TI**
*out* Y:**T**
*out* Y_fp32:**F**|2+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)| -|||1+|**F** = tensor(float)
**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)| +|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |ComplexMul|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ComplexMulConj|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc index 5b093a28f72aa..4b50b240bdea0 100644 --- a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc @@ -15,7 +15,6 @@ namespace onnxruntime::contrib { BranchwiseRMSNorm, kMSDomain, 1, T, kCpuExecutionProvider, \ KernelDefBuilder() \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("F", DataTypeImpl::GetTensorType()) \ .TypeConstraint("M", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}), \ @@ -48,11 +47,7 @@ template BranchwiseRMSNorm::BranchwiseRMSNorm(const OpKernelInfo& info) : LayerNormImpl(info, -1, info.GetAttrOrDefault("epsilon", 1e-5f), true), epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), - num_branches_(info.GetAttrOrDefault("num_branches", 0)), - use_onnx_arithmetic_(info.GetAttrOrDefault("use_onnx_arithmetic", 0) != 0) { - const auto mode = info.GetAttrOrDefault("use_onnx_arithmetic", 0); - ORT_ENFORCE(mode == 0 || mode == 1, "use_onnx_arithmetic must be 0 or 1"); -} + num_branches_(info.GetAttrOrDefault("num_branches", 0)) {} template Status BranchwiseRMSNorm::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, @@ -67,10 +62,6 @@ Status BranchwiseRMSNorm::PrePack(const Tensor& tensor, int input_idx, Alloca template Status BranchwiseRMSNorm::Compute(OpKernelContext* context) const { - const auto* numerics = context->Input(2); - ORT_RETURN_IF_NOT(numerics == nullptr || - (numerics->Shape().Size() == 1 && *numerics->Data() == 0), - "BranchwiseRMSNorm: nonzero prefill_numerics requires CUDA"); const auto* x = context->Input(0); const auto* scale = context->Input(1); hyper_connection::StreamShape params; @@ -80,10 +71,7 @@ Status BranchwiseRMSNorm::Compute(OpKernelContext* context) const { } auto* y = context->Output(0, x->Shape()); - auto* widened = context->OutputCount() > 1 ? context->Output(1, x->Shape()) : nullptr; - ORT_RETURN_IF_NOT(widened == nullptr || use_onnx_arithmetic_, - "Y_fp32 requires use_onnx_arithmetic=1"); - if (!use_onnx_arithmetic_ && scale != nullptr && + if (scale != nullptr && scale->Shape().Size() == params.hidden && scale->GetElementType() == x->GetElementType()) { TensorShapeVector grouped_dims(x->Shape().GetDims().begin(), @@ -110,20 +98,15 @@ Status BranchwiseRMSNorm::Compute(OpKernelContext* context) const { const float value = static_cast(x_data[offset + h]); sum_sq += value * value; } - const float rms = std::sqrt(sum_sq / static_cast(params.hidden) + epsilon_); - const float inv_rms = 1.0f / rms; + const float inv_rms = 1.0f / std::sqrt(sum_sq / static_cast(params.hidden) + epsilon_); for (int64_t h = 0; h < params.hidden; ++h) { const int64_t scale_index = scale != nullptr && scale->Shape().Size() == params.hidden ? h : (group % params.branches) * params.hidden + h; const float weight = scale == nullptr ? 1.0f : LoadFloat(*scale, scale_index); - const float value = static_cast(x_data[offset + h]); - y_data[offset + h] = static_cast( - (use_onnx_arithmetic_ ? value / rms : value * inv_rms) * weight); - if (widened != nullptr) { - widened->MutableData()[offset + h] = static_cast(y_data[offset + h]); - } + y_data[offset + h] = + static_cast(static_cast(x_data[offset + h]) * inv_rms * weight); } }, 0); diff --git a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h index bd58f86db058d..63284794f3f87 100644 --- a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h +++ b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h @@ -18,7 +18,6 @@ class BranchwiseRMSNorm final : public LayerNormImpl { private: float epsilon_; int64_t num_branches_; - bool use_onnx_arithmetic_; }; } // namespace onnxruntime::contrib diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc index 47aa148359aa1..ee21a71c2b4aa 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc @@ -9,31 +9,24 @@ #include "contrib_ops/cpu/hyper_connection_helper.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" -#include "core/providers/cuda/reduction/reduction_functions.h" namespace onnxruntime::contrib::cuda { using namespace onnxruntime::cuda; -#define REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(V, T) \ +#define REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ - BranchwiseRMSNorm, kMSDomain, V, T, kCudaExecutionProvider, \ + BranchwiseRMSNorm, kMSDomain, 1, T, kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("F", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ - .InputMemoryType(OrtMemTypeCPUInput, 2) \ .TypeConstraint("M", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}), \ BranchwiseRMSNorm); -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(1, float) -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(1, MLFloat16) -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(1, BFloat16) -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(2, float) -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(2, MLFloat16) -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(2, BFloat16) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(float) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(MLFloat16) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(BFloat16) #undef REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL @@ -41,83 +34,30 @@ template BranchwiseRMSNorm::BranchwiseRMSNorm(const OpKernelInfo& info) : CudaKernel(info), epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), - num_branches_(info.GetAttrOrDefault("num_branches", 0)), - use_onnx_arithmetic_(info.GetAttrOrDefault("use_onnx_arithmetic", 0) != 0), - input_start_(info.GetAttrOrDefault("input_start", 0)), - input_width_(info.GetAttrOrDefault("input_width", 0)) { - const auto mode = info.GetAttrOrDefault("use_onnx_arithmetic", 0); - ORT_ENFORCE(mode == 0 || mode == 1, "use_onnx_arithmetic must be 0 or 1"); - ORT_ENFORCE(input_start_ >= 0 && input_width_ >= 0, "input_start and input_width must be nonnegative"); - ORT_ENFORCE(input_width_ != 0 || input_start_ == 0, "input_start requires input_width"); - ORT_ENFORCE(input_width_ == 0 || (use_onnx_arithmetic_ && num_branches_ == 0), - "input_width requires use_onnx_arithmetic=1 and grouped input"); -} + num_branches_(info.GetAttrOrDefault("num_branches", 0)) {} template Status BranchwiseRMSNorm::ComputeInternal(OpKernelContext* context) const { using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; const auto* x = context->Input(0); const auto* scale = context->Input(1); - const auto* numerical_mode = context->Input(2); - bool dense_prefill = false; - if (numerical_mode != nullptr) { - ORT_RETURN_IF_NOT(numerical_mode->Shape().Size() == 1, - "prefill_numerics must contain one int32 value"); - const int32_t mode = *numerical_mode->Data(); - ORT_RETURN_IF_NOT(mode == 0 || mode == 1, "prefill_numerics must be 0 or 1"); - dense_prefill = mode == 1; - ORT_RETURN_IF_NOT(!dense_prefill || use_onnx_arithmetic_, - "Dense prefill numerical mode requires use_onnx_arithmetic=1"); - } - auto output_dims = x->Shape().AsShapeVector(); - ORT_RETURN_IF_NOT(!output_dims.empty(), "BranchwiseRMSNorm requires a non-scalar input"); - const int64_t input_stride = output_dims.back(); - if (input_width_ != 0) { - ORT_RETURN_IF_NOT(input_start_ <= input_stride && input_width_ <= input_stride - input_start_, - "input slice exceeds the last dimension"); - output_dims.back() = input_width_; - } - const TensorShape output_shape(output_dims); hyper_connection::StreamShape params; - ORT_RETURN_IF_ERROR(hyper_connection::ResolveStreamShape(output_shape, num_branches_, params)); + ORT_RETURN_IF_ERROR(hyper_connection::ResolveStreamShape(x->Shape(), num_branches_, params)); if (scale != nullptr) { ORT_RETURN_IF_ERROR(hyper_connection::ValidateScale(scale->Shape(), params)); } ORT_RETURN_IF_NOT(params.branches <= std::numeric_limits::max() && - params.hidden <= std::numeric_limits::max() && - (input_width_ == 0 || input_stride <= std::numeric_limits::max()), + params.hidden <= std::numeric_limits::max(), "branch or hidden dimension is too large for CUDA"); - const int64_t groups = output_shape.Size() / params.hidden; - auto* y = context->Output(0, output_shape); - auto* widened = context->OutputCount() > 1 ? context->Output(1, output_shape) : nullptr; - ORT_RETURN_IF_NOT(widened == nullptr || use_onnx_arithmetic_, - "Y_fp32 requires use_onnx_arithmetic=1"); + const int64_t groups = x->Shape().Size() / params.hidden; + auto* y = context->Output(0, x->Shape()); if (groups == 0) { return Status::OK(); } const auto* x_data = reinterpret_cast(x->Data()); auto* y_data = reinterpret_cast(y->MutableData()); - if (use_onnx_arithmetic_) { - ORT_RETURN_IF_NOT(groups <= std::numeric_limits::max(), - "ONNX arithmetic requires the row count to fit in int32"); - const auto buffer_size = compute_reduce_matrix_columns_buffer_size( - static_cast(groups), static_cast(params.hidden)); - const bool deterministic_compute = context->GetUseDeterministicCompute(); - auto squares = GetScratchBuffer(groups == 1 && !deterministic_compute ? output_shape.Size() : 0, - GetComputeStream(context)); - auto sums = GetScratchBuffer(groups, GetComputeStream(context)); - auto buffer = GetScratchBuffer(buffer_size, GetComputeStream(context)); - return LaunchStandardBranchwiseRMSNorm( - Stream(context), x_data + input_start_, scale == nullptr ? nullptr : scale->DataRaw(), - scale == nullptr ? x->GetElementType() : scale->GetElementType(), y_data, - static_cast(groups), static_cast(params.branches), static_cast(params.hidden), - scale == nullptr || scale->Shape().Size() == params.hidden, epsilon_, - widened == nullptr ? nullptr : widened->MutableData(), - squares.get(), sums.get(), buffer.get(), buffer_size, - static_cast(input_width_ == 0 ? params.hidden : input_stride), deterministic_compute, dense_prefill); - } return LaunchMixedScaleBranchwiseRMSNorm( Stream(context), x_data, scale == nullptr ? nullptr : scale->DataRaw(), scale == nullptr ? x->GetElementType() : scale->GetElementType(), y_data, diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h index c26a8d47b41ff..959fa12b4a46a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h @@ -16,9 +16,6 @@ class BranchwiseRMSNorm final : public onnxruntime::cuda::CudaKernel { private: float epsilon_; int64_t num_branches_; - bool use_onnx_arithmetic_; - int64_t input_start_; - int64_t input_width_; }; } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu index c481e05cc8db1..a56e55ce21331 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu @@ -7,12 +7,10 @@ #include #include #include -#include #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" #include "core/providers/cuda/cuda_common.h" -#include "core/providers/cuda/reduction/reduction_functions.h" namespace onnxruntime::contrib::cuda { namespace { @@ -31,35 +29,6 @@ __device__ __forceinline__ float LoadScale(const void* data, int64_t index, int return __bfloat162float(static_cast(data)[index]); } -template -__global__ void BranchwiseSquareKernel(const T* x, float* squares, int64_t count) { - const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (index < count) { - const float value = to_float(x[index]); - squares[index] = __fmul_rn(value, value); - } -} - -template -__global__ void BranchwiseStandardEpilogueKernel( - const T* x, const void* scale, int scale_type, const float* sums, T* y, float* widened, - int64_t count, int branches, int hidden, bool shared_scale, float epsilon, int input_row_stride, bool reciprocal_multiply) { - const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (index < count) { - const int64_t group = index / hidden; - const int64_t scale_index = shared_scale ? index % hidden : index % (static_cast(branches) * hidden); - const float weight = scale == nullptr ? 1.0f : LoadScale(scale, scale_index, scale_type); - const float mean = sums[group] / static_cast(hidden); - const float rms = sqrtf(__fadd_rn(mean, epsilon)); - const int64_t source_index = Packed ? group * input_row_stride + index % hidden : index; - const float value = to_float(x[source_index]); - const float normalized = reciprocal_multiply ? __fmul_rn(value, __fdiv_rn(1.0f, rms)) : value / rms; - const T rounded = from_float(__fmul_rn(normalized, weight)); - y[index] = rounded; - if (widened != nullptr) widened[index] = to_float(rounded); - } -} - template __global__ void MixedScaleBranchwiseRMSNormKernel(const T* x, const void* scale, int scale_type, T* y, int64_t groups, int branches, int hidden, @@ -108,39 +77,6 @@ __global__ void MixedScaleBranchwiseRMSNormKernel(const T* x, const void* scale, } // namespace -template -Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const void* scale, - int scale_type, T* y, int groups, int branches, - int hidden, bool shared_scale, float epsilon, - float* widened, float* squares, float* sums, void* reduction_buffer, - size_t reduction_buffer_size, int input_row_stride, - bool deterministic_compute, bool reciprocal_multiply) { - const int64_t count = static_cast(groups) * hidden; - const int64_t blocks = (count + kThreads - 1) / kThreads; - ORT_RETURN_IF_NOT(blocks <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); - // Standard ReduceMean removes singleton dimensions and prefers row reduction for a scalar result. - // Row reduction uses floating-point atomics; column reduction combines indexed partials in a fixed order. - if (groups == 1 && !deterministic_compute) { - BranchwiseSquareKernel<<(blocks), kThreads, 0, stream>>>(x, squares, count); - ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); - ORT_RETURN_IF_ERROR((onnxruntime::cuda::reduce_matrix_rows( - stream, squares, sums, hidden, 1))); - } else { - using ReduceT = std::conditional_t, BFloat16, T>; - ORT_RETURN_IF_ERROR((onnxruntime::cuda::reduce_square_matrix_columns( - stream, reinterpret_cast(x), sums, groups, hidden, - reduction_buffer, reduction_buffer_size, input_row_stride))); - } - if (input_row_stride == hidden) { - BranchwiseStandardEpilogueKernel<<(blocks), kThreads, 0, stream>>>( - x, scale, scale_type, sums, y, widened, count, branches, hidden, shared_scale, epsilon, input_row_stride, reciprocal_multiply); - } else { - BranchwiseStandardEpilogueKernel<<(blocks), kThreads, 0, stream>>>( - x, scale, scale_type, sums, y, widened, count, branches, hidden, shared_scale, epsilon, input_row_stride, reciprocal_multiply); - } - return CUDA_CALL(cudaGetLastError()); -} - template Status LaunchMixedScaleBranchwiseRMSNorm(cudaStream_t stream, const T* x, const void* scale, int scale_type, T* y, int64_t groups, int branches, @@ -171,12 +107,4 @@ INSTANTIATE(__nv_bfloat16) #undef INSTANTIATE -#define INSTANTIATE_STANDARD(T) \ - template Status LaunchStandardBranchwiseRMSNorm(cudaStream_t, const T*, const void*, int, T*, \ - int, int, int, bool, float, float*, float*, float*, void*, size_t, int, bool, bool); -INSTANTIATE_STANDARD(float) -INSTANTIATE_STANDARD(half) -INSTANTIATE_STANDARD(__nv_bfloat16) -#undef INSTANTIATE_STANDARD - } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h index c878e8389b6a2..549de28e62bff 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h @@ -14,12 +14,4 @@ Status LaunchMixedScaleBranchwiseRMSNorm(cudaStream_t stream, const T* x, const int scale_type, T* y, int64_t groups, int branches, int hidden, bool shared_scale, float epsilon); -template -Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const void* scale, - int scale_type, T* y, int groups, int branches, - int hidden, bool shared_scale, float epsilon, - float* widened, float* squares, float* sums, void* reduction_buffer, - size_t reduction_buffer_size, int input_row_stride, - bool deterministic_compute, bool reciprocal_multiply = false); - } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc index 369489f585baa..810016b96529c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc @@ -4,13 +4,8 @@ #include "contrib_ops/cuda/bert/rms_norm_rotary.h" #include +#include "core/common/safeint.h" #include "core/providers/cuda/cuda_common.h" -#include "core/providers/cuda/cuda_type_conversion.h" -#include "core/providers/cuda/shared_inc/fpgeneric.h" -#include "contrib_ops/cuda/bert/branchwise_rms_norm_impl.h" -#include "core/providers/cuda/llm/rotary_embedding_impl.h" -#include "core/providers/cuda/reduction/reduction_functions.h" -#include "contrib_ops/cuda/bert/rotary_embedding_impl.h" namespace onnxruntime::contrib::cuda { ONNX_OPERATOR_KERNEL_EX( @@ -50,7 +45,7 @@ Status RMSNormRotary::ComputeInternal(OpKernelContext* context) const { 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 && input_start_ <= shape[3] && width <= shape[3] - input_start_ && + 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(); @@ -72,44 +67,19 @@ Status RMSNormRotary::ComputeInternal(OpKernelContext* context) const { 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), batch = SafeInt(output_batch); + 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_; - if (dim == 256 && groups > 1 && (dense || fp32_rotary_)) { - 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, rotary, cache_rows, scalar, epsilon_, dense, fp32_rotary_); - } - using CudaT = onnxruntime::cuda::OrtToCudaType::type; - auto prepared = GetScratchBuffer(SafeInt(norm_shape.Size()), GetComputeStream(context)); - auto squares = GetScratchBuffer(groups == 1 ? width : 0, GetComputeStream(context)); - auto sums = GetScratchBuffer(groups, GetComputeStream(context)); - const auto buffer_size = onnxruntime::cuda::compute_reduce_matrix_columns_buffer_size(groups, dim); - auto buffer = GetScratchBuffer(buffer_size, GetComputeStream(context)); - ORT_RETURN_IF_ERROR(LaunchStandardBranchwiseRMSNorm( - Stream(context), reinterpret_cast(x), scale->Data(), - ONNX_NAMESPACE::TensorProto::FLOAT, prepared.get(), groups, heads, dim, true, epsilon_, - normalized ? normalized->MutableData() : nullptr, squares.get(), sums.get(), - buffer.get(), buffer_size, stride, context->GetUseDeterministicCompute(), dense)); - const auto launch_rotary = [&]() { - return LaunchRotaryEmbeddingKernel( - Stream(context), reinterpret_cast(output->MutableData()), - reinterpret_cast(prepared.get()), - positions->Data(), nullptr, - reinterpret_cast(cosine->Data()), - reinterpret_cast(sine->Data()), batch, sequence, heads, dim, rotary, - cache_rows, scalar ? 0 : 1, false, GetDeviceProp().maxThreadsPerBlock, false, - !dense && fp32_rotary_); - }; - if (dense) return launch_rotary.template operator()(); - return launch_rotary.template operator()::MappedType>(); + 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 index 42acb0ff84197..6014e47983be1 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu @@ -9,33 +9,29 @@ 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 rotary_dim, + 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; - constexpr int width = 256; __shared__ float partial[2]; - __shared__ __nv_bfloat16 values[width]; + __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 v[4]; float sum = 0.f; -#pragma unroll - for (int i = 0; i < 4; ++i) v[i] = __bfloat162float(x[tid + i * threads]); - // Match reduce_square_matrix_columns' four strided values and two-warp shuffle tree. -#pragma unroll - for (int i = 0; i < 4; ++i) sum += v[i] * v[i]; + 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 / float(width), epsilon)); -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int column = tid + i * threads; - const float norm = dense_prefill ? __fmul_rn(v[i], __fdiv_rn(1.f, rms)) : v[i] / rms; + 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); @@ -46,9 +42,7 @@ __global__ void RMSNormRotaryKernel( : base >= 0 && base < cache_rows; const int64_t position = valid && scalar_position ? base + token % sequence : base; const int half = rotary_dim / 2; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int column = tid + i * threads; + for (int column = tid; column < width; column += threads) { auto value = values[column]; if (column < rotary_dim && valid) { const int partner = (column + half) % rotary_dim; @@ -71,17 +65,17 @@ __global__ void RMSNormRotaryKernel( 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 rotary_dim, int cache_rows, bool scalar_position, + 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 && input_stride >= 256 && - rotary_dim > 0 && rotary_dim <= 256 && rotary_dim % 2 == 0 && cache_rows > 0 && - (dense_prefill || fp32_rotary), - "RMSNormRotary fast path requires valid dimensions and width 256"); + 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, - rotary_dim, cache_rows, scalar_position, epsilon, dense_prefill, fp32_rotary); + width, rotary_dim, cache_rows, scalar_position, epsilon, dense_prefill, fp32_rotary); return CUDA_CALL(cudaGetLastError()); } diff --git a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h index f6220658b8413..6b71093b2b3e8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h +++ b/onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h @@ -7,7 +7,7 @@ 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 rotary_dim, int cache_rows, bool scalar_position, + 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 { diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu index 4415ecc5fef4f..47c46b73f7b83 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu @@ -10,7 +10,6 @@ Kernel implementation for rotary embeddings. #include "contrib_ops/cuda/bert/rotary_embedding_impl.h" #include -#include #include "core/common/safeint.h" #include "core/providers/cuda/cu_inc/common.cuh" @@ -21,7 +20,7 @@ namespace onnxruntime { namespace contrib { namespace cuda { -template +template __global__ void RotaryEmbeddingBSNH(T* output, // BxSxNxH const T* input, // BxSxNxH const T* cos_cache, // Mx(H/2) @@ -129,14 +128,7 @@ __global__ void RotaryEmbeddingBSNH(T* output, // BxSxNx } // Use values from shared memory - if constexpr (use_fp32_arithmetic) { - const float a = static_cast(use_smem ? smem[i] : input_data[i]); - const float b = static_cast(use_smem ? smem[j] : input_data[j]); - const float cosine = static_cast(cos_data[cache_idx]); - const float sine = static_cast(sin_data[cache_idx]); - const float signed_b = sign < T(0) ? -b : b; - output_data[i] = T(__fadd_rn(__fmul_rn(a, cosine), __fmul_rn(signed_b, sine))); - } else if constexpr (use_smem) { + if constexpr (use_smem) { output_data[i] = smem[i] * cos_data[cache_idx] + sign * smem[j] * sin_data[cache_idx]; } else { output_data[i] = input_data[i] * cos_data[cache_idx] + sign * input_data[j] * sin_data[cache_idx]; @@ -150,8 +142,7 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, - const int max_threads_per_block, const bool is_input_bnsh_format, - const bool use_fp32_arithmetic) { + const int max_threads_per_block, const bool is_input_bnsh_format) { int sequence_stride; int batch_stride; int heads_stride; @@ -177,7 +168,7 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu rotary_embedding_dim, max_sequence_length, position_ids_format, interleaved, max_threads_per_block, - in_strides, out_strides, use_fp32_arithmetic); + in_strides, out_strides); } template @@ -188,8 +179,8 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, const int max_threads_per_block, - int4 in_strides, int4 out_strides, // strides in bnsh coord - const bool use_fp32_arithmetic) { + int4 in_strides, int4 out_strides // strides in bnsh coord +) { // Note: Current implementation assumes head_size <= max_threads_per_block // because head_size is currently large for LLaMA-2. For smaller head_size // and num_heads values, we can create a block as `block(num_heads, head_size, 1)` @@ -205,25 +196,20 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu assert(head_size <= max_threads_per_block); - auto launch = [&](auto mode) { - constexpr bool fp32 = decltype(mode)::value; - if (output == input) { - size_t smem_size = head_size * sizeof(T); - RotaryEmbeddingBSNH<<>>( - output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, - num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, - interleaved, in_strides, out_strides); - } else { - RotaryEmbeddingBSNH<<>>( - output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, - num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, - interleaved, in_strides, out_strides); - } - }; - if (use_fp32_arithmetic) - launch(std::true_type{}); - else - launch(std::false_type{}); + if (output == input) { + // In-place operation: use shared memory to avoid read-after-write hazards + size_t smem_size = head_size * sizeof(T); + RotaryEmbeddingBSNH<<>>( + output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, + num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, + interleaved, in_strides, out_strides); + } else { + // Separate buffers: no shared memory needed + RotaryEmbeddingBSNH<<>>( + output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, + num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, + interleaved, in_strides, out_strides); + } return CUDA_CALL(cudaGetLastError()); } @@ -234,8 +220,7 @@ template Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, float* o const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, - const int max_threads_per_block, const bool is_input_bnsh_format, - const bool use_fp32_arithmetic); + const int max_threads_per_block, const bool is_input_bnsh_format); template Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, half* output, const half* input, const int64_t* position_ids, const int* past_sequence_lengths, const half* cos_cache, @@ -243,19 +228,14 @@ template Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, half* out const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, - const int max_threads_per_block, const bool is_input_bnsh_format, - const bool use_fp32_arithmetic); - -template Status LaunchRotaryEmbeddingKernel<__nv_bfloat16>( - cudaStream_t, __nv_bfloat16*, const __nv_bfloat16*, const int64_t*, const int*, - const __nv_bfloat16*, const __nv_bfloat16*, int, int, int, int, int, int, int, bool, int, bool, bool); + const int max_threads_per_block, const bool is_input_bnsh_format); template Status LaunchRotaryEmbeddingKernel( cudaStream_t stream, BFloat16* output, const BFloat16* input, const int64_t* position_ids, const int* past_sequence_lengths, const BFloat16* cos_cache, const BFloat16* sin_cache, const int batch_size, const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, const int max_threads_per_block, - const bool is_input_bnsh_format, const bool use_fp32_arithmetic); + const bool is_input_bnsh_format); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h index 4422572fc271d..2d81e6b4067af 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h @@ -27,8 +27,7 @@ Status LaunchRotaryEmbeddingKernel( const int position_ids_format, const bool interleaved, const int max_threads_per_block, - const bool is_input_bnsh_format, - const bool use_fp32_arithmetic = false); + const bool is_input_bnsh_format); template Status LaunchRotaryEmbeddingKernel( @@ -49,8 +48,7 @@ Status LaunchRotaryEmbeddingKernel( const bool interleaved, const int max_threads_per_block, int4 in_strides, - int4 out_strides, - const bool use_fp32_arithmetic = false); + int4 out_strides); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index d1b0c3c39bea1..5d349ddf1bade 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -180,9 +180,6 @@ 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); -class CUDA_MS_OP_TYPED_CLASS_NAME(2, float, BranchwiseRMSNorm); -class CUDA_MS_OP_TYPED_CLASS_NAME(2, MLFloat16, BranchwiseRMSNorm); -class CUDA_MS_OP_TYPED_CLASS_NAME(2, BFloat16, BranchwiseRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, ScaledSiLU); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, ScaledSiLU); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, ScaledSiLU); @@ -551,9 +548,6 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc index cb06ad18f3754..44747120cc2bc 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc @@ -15,7 +15,6 @@ ONNX_OPERATOR_KERNEL_EX( BranchwiseRMSNorm, kMSDomain, 1, kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", WebGpuSupportedFloatTypes()) - .TypeConstraint("F", DataTypeImpl::GetTensorType()) .TypeConstraint("M", WebGpuSupportedFloatTypes()), BranchwiseRMSNorm); @@ -62,17 +61,9 @@ Status BranchwiseRMSNormProgram::GenerateShaderCode(ShaderHelper& shader) const BranchwiseRMSNorm::BranchwiseRMSNorm(const OpKernelInfo& info) : WebGpuKernel(info), epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), - num_branches_(info.GetAttrOrDefault("num_branches", 0)) { - ORT_ENFORCE(info.GetAttrOrDefault("use_onnx_arithmetic", 0) == 0, - "WebGPU BranchwiseRMSNorm does not support use_onnx_arithmetic=1"); -} + num_branches_(info.GetAttrOrDefault("num_branches", 0)) {} Status BranchwiseRMSNorm::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { - const auto* numerics = context.Input(2); - ORT_RETURN_IF_NOT(numerics == nullptr || - (numerics->IsDataType() && numerics->Shape().Size() == 1 && - numerics->Data()[0] == 0), - "BranchwiseRMSNorm: nonzero prefill_numerics requires CUDA"); const auto* x = context.Input(0); const auto* scale = context.Input(1); hyper_connection::StreamShape params; diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 028bd81694fcf..56d57cec1d700 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -5851,94 +5851,31 @@ Applies RMS normalization independently to each branch. X may use grouped shape (..., C, H), or flattened shape (..., C * H) when num_branches is specified. The optional scale may have shape (C * H), (C, H), or (H). Arithmetic is performed in float32 and the result is converted to T. -With use_onnx_arithmetic=1, compute FP32 square, mean, epsilon addition, square -root, division, and scale multiplication with separate rounding boundaries. -CUDA reuses the standard fast last-axis ReduceMean summation order by default. -Deterministic CUDA sessions instead use a fixed-order reduction, which may -round differently from the default reduction or deterministic ReduceMean. -This mode does not substitute reciprocal square root. -Optional Y_fp32 losslessly widens the already T-rounded Y to float32; it does -not bypass the output rounding boundary and requires use_onnx_arithmetic=1. -WebGPU currently supports only the default mode without Y_fp32. )DOC"; -static OpSchema BranchwiseRMSNormSchema() { - return OpSchema() - .SetDoc(BranchwiseRMSNorm_ver1_doc) - .Attr("epsilon", "Epsilon added before reciprocal square root.", - AttributeProto::FLOAT, 1e-5f) - .Attr("num_branches", - "Number of branches for flattened input. Omit or set to zero for grouped input.", - AttributeProto::INT, static_cast(0)) - .Attr("use_onnx_arithmetic", - "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. " - "Supported values are 0 (default) and 1.", - AttributeProto::INT, static_cast(0)) - .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") - .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", - "M", OpSchema::Optional) - .Input(2, "prefill_numerics", "Optional CPU int32 scalar 0/1; one enables dense prefill arithmetic.", "TI", OpSchema::Optional) - .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") - .Output(1, "Y_fp32", - "Optional float32 widening of the already T-rounded Y. " - "Requires use_onnx_arithmetic=1.", - "F", OpSchema::Optional) - .TypeConstraint("T", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain input and output to floating-point tensors.") - .TypeConstraint("M", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain scale to floating-point tensors.") - .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.") - .TypeConstraint("TI", {"tensor(int32)"}, "CPU numerical control.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - const auto start = getAttribute(ctx, "input_start", 0); - const auto width = getAttribute(ctx, "input_width", 0); - if (start < 0 || width < 0 || (width == 0 && start != 0)) { - fail_shape_inference("input_start and input_width must describe a nonnegative slice"); - } - if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 || - getAttribute(ctx, "num_branches", 0) != 0)) { - fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input"); - } - if (hasInputShape(ctx, 0)) { - auto shape = getInputShape(ctx, 0); - if (width != 0) { - if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two"); - auto* last = shape.mutable_dim(shape.dim_size() - 1); - if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) { - fail_shape_inference("input slice exceeds the last dimension"); - } - last->set_dim_value(width); - } - updateOutputShape(ctx, 0, shape); - } - if (ctx.getNumOutputs() > 1) { - if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) { - fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1"); - } - updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - if (hasInputShape(ctx, 0)) { - auto shape = getInputShape(ctx, 0); - if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width); - updateOutputShape(ctx, 1, shape); - } - } - }); -} - -ONNX_MS_OPERATOR_SET_SCHEMA(BranchwiseRMSNorm, 1, BranchwiseRMSNormSchema()); ONNX_MS_OPERATOR_SET_SCHEMA( - BranchwiseRMSNorm, 2, - BranchwiseRMSNormSchema() - .Attr("input_start", "CUDA-only starting offset in the last input dimension.", + BranchwiseRMSNorm, 1, + OpSchema() + .SetDoc(BranchwiseRMSNorm_ver1_doc) + .Attr("epsilon", "Epsilon added before reciprocal square root.", + AttributeProto::FLOAT, 1e-5f) + .Attr("num_branches", + "Number of branches for flattened input. Omit or set to zero for grouped input.", AttributeProto::INT, static_cast(0)) - .Attr("input_width", - "CUDA-only width selected from the last input dimension before normalization. " - "Zero retains the full input. Requires grouped input and use_onnx_arithmetic=1. " - "Outputs retain the leading dimensions and use this last dimension.", - AttributeProto::INT, static_cast(0))); + .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") + .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", + "M", OpSchema::Optional) + .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain input and output to floating-point tensors.") + .TypeConstraint("M", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain scale to floating-point tensors.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateShapeFromInputToOutput(ctx, 0, 0); + })); constexpr const char* ScaledSiLU_ver1_doc = R"DOC( Computes SiLU after scaling, with explicit T rounding: diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index f54b308b5726e..d45e3d4f65b28 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -92,7 +92,6 @@ 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, 2, BranchwiseRMSNorm); 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); @@ -230,7 +229,6 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); - fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index 75024a420d3a9..60ffd772c97f1 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -275,7 +275,7 @@ __device__ void reduce_all( template __global__ void reduce_matrix_columns_kernel( const int num_rows, const int num_cols, const TIn* const input, TOut* const output, - TBuf* const block_reductions_buffer, int* const block_done_counts_buffer, const int input_row_stride) { + TBuf* const block_reductions_buffer, int* const block_done_counts_buffer) { const int num_blocks_in_grid_row = gridDim.x; const int row_id_in_grid = blockIdx.y; const int num_grid_rows = gridDim.y; @@ -283,7 +283,7 @@ __global__ void reduce_matrix_columns_kernel( // one row per iteration // row_id is int64_t to avoid int overflow in offset calculations for (int64_t row_id = row_id_in_grid; row_id < num_rows; row_id += num_grid_rows) { - const TIn* const row_data = input + row_id * input_row_stride; + const TIn* const row_data = input + row_id * num_cols; TOut* const row_output = output + row_id; TBuf* const row_block_reductions_buffer = block_reductions_buffer + row_id * num_blocks_in_grid_row; int* const row_block_done_counts_buffer = block_done_counts_buffer + row_id; @@ -296,10 +296,8 @@ __global__ void reduce_matrix_columns_kernel( template Status call_reduce_matrix_columns( - cudaStream_t stream, const TIn* input, TOut* output, const int num_rows, const int num_cols, - void* buffer, size_t buffer_size, int input_row_stride = 0) { + cudaStream_t stream, const TIn* input, TOut* output, const int num_rows, const int num_cols, void* buffer, size_t buffer_size) { ORT_ENFORCE(num_rows >= 0 && num_cols >= 0); - ORT_ENFORCE(input_row_stride == 0 || input_row_stride >= num_cols); using TBuf = AccumulationType_t; @@ -321,8 +319,7 @@ Status call_reduce_matrix_columns( const int shared_mem_size = sizeof(TBuf) * block_dim.x * block_dim.y / GPU_WARP_SIZE_HOST; reduce_matrix_columns_kernel <<>>( - num_rows, num_cols, input, output, block_reductions_buffer, block_done_counts_buffer, - input_row_stride == 0 ? num_cols : input_row_stride); + num_rows, num_cols, input, output, block_reductions_buffer, block_done_counts_buffer); return Status::OK(); } @@ -531,26 +528,6 @@ INSTANTIATE_REDUCE_MATRIX_COLUMNS(double); INSTANTIATE_REDUCE_MATRIX_COLUMNS(BFloat16); #undef INSTANTIATE_REDUCE_MATRIX_COLUMNS -struct RoundedSquare { - __device__ __forceinline__ float operator()(float value) const { - return __fmul_rn(value, value); - } -}; - -template -Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, - int m, int n, void* buffer, size_t buffer_size, int input_row_stride) { - return detail::call_reduce_matrix_columns( - stream, input, output, m, n, buffer, buffer_size, input_row_stride); -} - -#define INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(T) \ - template Status reduce_square_matrix_columns(cudaStream_t, const T*, float*, int, int, void*, size_t, int) -INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(float); -INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(half); -INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(BFloat16); -#undef INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS - namespace detail { constexpr int kMaxReduceRank = 16; diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index 9c877c0743b31..b772876db6395 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h @@ -106,11 +106,6 @@ Status reduce_matrix_rows(cudaStream_t stream, const TIn* input, TOut* output, i template Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output, int m, int n, void* buffer, size_t buffer_size); -/** Sum FP32-rounded squares across columns, preserving reduce_matrix_columns' summation order. */ -template -Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, - int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0); - /** * Computes ReduceSum for an arbitrary set of axes. This is the general fallback * for axis layouts that cannot be flattened into one of the optimized matrix diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc index dab95a3c21d8b..bc8009b05c9e7 100644 --- a/onnxruntime/core/session/environment.cc +++ b/onnxruntime/core/session/environment.cc @@ -296,7 +296,7 @@ Status Environment::Initialize(std::unique_ptr logging_ auto& domainToVersionRangeInstance = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance(); if (domainToVersionRangeInstance.Map().find(onnxruntime::kMSDomain) == domainToVersionRangeInstance.Map().end()) { // External shared providers may have already added kMSDomain - domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 2); + domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 1); } domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSExperimentalDomain, 1, 1); domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSNchwcDomain, 1, 1); diff --git a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc index 46401997bd4ff..4a24c77b1d0f0 100644 --- a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc +++ b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc @@ -2,9 +2,7 @@ // Licensed under the MIT License. #include -#include #include -#include #include #include @@ -12,7 +10,6 @@ #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" -#include "core/framework/session_options.h" namespace onnxruntime::test { namespace { @@ -145,371 +142,6 @@ TEST(HyperConnectionOpsTest, BranchwiseRMSNormMixedScaleType) { }); } -class StandardRMSChainTester : public OpTester { - public: - StandardRMSChainTester(int output_type, float epsilon, int64_t start = 0, int64_t width = 0) - : OpTester("StandardRMSChain", 17, kOnnxDomain, false), - output_type_(output_type), - epsilon_(epsilon), - start_(start), - width_(width) {} - - void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, - std::vector>&) override { - auto value = [&](const char* name) { return &graph.GetOrCreateNodeArg(name, nullptr); }; - NodeArg* source = inputs[0]; - if (width_ != 0) { - auto constant = [&](const char* name, int64_t data) { - ONNX_NAMESPACE::TensorProto tensor; - tensor.set_name(name); - tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); - tensor.add_dims(1); - tensor.add_int64_data(data); - graph.AddInitializedTensor(tensor); - return value(name); - }; - source = value("selected"); - graph.AddNode("slice", "Slice", "", - {inputs[0], constant("starts", start_), constant("ends", start_ + width_), - constant("axes", -1)}, - {source}); - } - auto* x = value("float_x"); - graph.AddNode("cast", "Cast", "", {source}, {x}) - .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); - auto* square = value("square"); - graph.AddNode("square", "Mul", "", {x, x}, {square}); - auto* mean = value("mean"); - auto& reduce = graph.AddNode("mean", "ReduceMean", "", {square}, {mean}); - reduce.AddAttribute("axes", std::vector{-1}); - reduce.AddAttribute("keepdims", int64_t{1}); - ONNX_NAMESPACE::TensorProto epsilon; - epsilon.set_name("epsilon"); - epsilon.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - epsilon.add_float_data(epsilon_); - graph.AddInitializedTensor(epsilon); - auto* adjusted = value("adjusted"); - graph.AddNode("epsilon", "Add", "", {mean, value("epsilon")}, {adjusted}); - auto* root = value("root"); - graph.AddNode("root", "Sqrt", "", {adjusted}, {root}); - auto* normalized = value("normalized"); - graph.AddNode("divide", "Div", "", {x, root}, {normalized}); - auto* scaled = value("scaled"); - graph.AddNode("scale", "Mul", "", {normalized, inputs[1]}, {scaled}); - graph.AddNode("output", "Cast", "", {scaled}, {outputs[0]}) - .AddAttribute("to", static_cast(output_type_)); - if (outputs.size() > 1) { - graph.AddNode("widen", "Cast", "", {outputs[0]}, {outputs[1]}) - .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); - } - graph.SetOutputs(outputs); - } - - private: - int output_type_; - float epsilon_; - int64_t start_; - int64_t width_; -}; - -template -void TestStandardRMSArithmetic(int type, bool packed = false) { - if (!DefaultCudaExecutionProvider()) { - GTEST_SKIP() << "CUDA is required for exact standard RMS arithmetic"; - } - for (int width : {1, 31, 128, 256, 513, 1024, 4096}) { - for (int rows : (packed ? std::vector{1, 2, 9, 48} : std::vector{1, 2, 9})) { - for (int start : (packed ? std::vector{0, 3} : std::vector{0})) { - SCOPED_TRACE("width=" + std::to_string(width) + " rows=" + std::to_string(rows)); - const int stride = packed ? width * 2 + 7 : width; - std::vector input(rows * stride); - std::vector scale(width); - for (int i = 0; i < rows * stride; ++i) { - input[i] = T(std::sin(static_cast(i) * 0.71f) * (1 + i % 17)); - } - for (int i = 0; i < width; ++i) scale[i] = 0.5f + static_cast(i % 23) / 17.0f; - constexpr float epsilon = 1e-6f; - StandardRMSChainTester reference(type, epsilon, start, packed ? width : 0); - OpTester fused("BranchwiseRMSNorm", packed ? 2 : 1, kMSDomain, false); - if (packed) { - fused.AddAttribute("input_start", start); - fused.AddAttribute("input_width", width); - } - fused.AddAttribute("use_onnx_arithmetic", 1); - fused.AddAttribute("epsilon", epsilon); - std::vector expected; - for (OpTester* tester : {static_cast(&reference), &fused}) { - const std::vector input_shape = rows == 48 ? std::vector{1, 2, 24, stride} - : std::vector{1, rows, stride}; - auto output_shape = input_shape; - output_shape.back() = width; - tester->AddInput("X", input_shape, input); - tester->AddInput("scale", {width}, scale); - tester->AddOutput("Y", output_shape, std::vector(rows * width)); - tester->AddOutput("Y_fp32", output_shape, std::vector(rows * width)); - SessionOptions options; - options.graph_optimization_level = TransformerLevel::Default; - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - tester->Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); - auto fetches = tester->GetFetches(); - ASSERT_EQ(fetches.size(), 2); - if (expected.empty()) { - expected = std::move(fetches); - } else { - for (size_t output = 0; output < fetches.size(); ++output) { - const auto& before = expected[output].Get(); - const auto& after = fetches[output].Get(); - ASSERT_EQ(before.SizeInBytes(), after.SizeInBytes()); - EXPECT_EQ(std::memcmp(before.DataRaw(), after.DataRaw(), before.SizeInBytes()), 0) - << "output=" << output; - } - } - } - } - } - } -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardFloatExact) { - TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardFloat16Exact) { - TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardBFloat16Exact) { - TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedFloatExact) { - TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, true); -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedFloat16Exact) { - TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, true); -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedBFloat16Exact) { - TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16, true); -} - -template -void TestStandardRMSBranchScales(int type) { - if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for exact branch scales"; - constexpr int outer = 2; - constexpr int branches = 3; - constexpr float epsilon = 1e-6f; - for (int width : {31, 128, 513}) { - for (bool flattened_scale : {false, true}) { - SCOPED_TRACE(width); - SCOPED_TRACE(flattened_scale); - std::vector input(outer * branches * width); - std::vector scale(branches * width); - std::vector float_scale(scale.size()); - for (size_t i = 0; i < input.size(); ++i) input[i] = T(std::sin(static_cast(i) * 0.71f) * 9.0f); - for (int branch = 0; branch < branches; ++branch) { - for (int h = 0; h < width; ++h) { - const int index = branch * width + h; - scale[index] = ScaleT((branch + 1) * 0.7f + (h % 17) * 0.013f); - if constexpr (std::is_same_v) { - float_scale[index] = scale[index]; - } else { - float_scale[index] = scale[index].ToFloat(); - } - } - } - const std::vector shape{outer, branches, width}; - StandardRMSChainTester reference(type, epsilon); - reference.AddInput("X", shape, input); - reference.AddInput("scale", {branches, width}, float_scale); - reference.AddOutput("Y", shape, std::vector(input.size())); - reference.AddOutput("Y_fp32", shape, std::vector(input.size())); - OpTester fused("BranchwiseRMSNorm", 2, kMSDomain, false); - fused.AddAttribute("use_onnx_arithmetic", 1); - fused.AddAttribute("epsilon", epsilon); - fused.AddInput("X", shape, input); - fused.AddInput("scale", flattened_scale ? std::vector{branches * width} : std::vector{branches, width}, - scale); - fused.AddOutput("Y", shape, std::vector(input.size())); - fused.AddOutput("Y_fp32", shape, std::vector(input.size())); - std::vector expected; - for (OpTester* tester : {static_cast(&reference), &fused}) { - SessionOptions options; - options.graph_optimization_level = TransformerLevel::Default; - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - tester->Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); - const auto& fetches = tester->GetFetches(); - ASSERT_EQ(fetches.size(), 2u); - if (expected.empty()) { - expected = fetches; - } else { - for (size_t output = 0; output < fetches.size(); ++output) { - const auto& before = expected[output].Get(); - const auto& after = fetches[output].Get(); - ASSERT_EQ(before.SizeInBytes(), after.SizeInBytes()); - EXPECT_EQ(std::memcmp(before.DataRaw(), after.DataRaw(), before.SizeInBytes()), 0) << output; - } - } - } - } - } -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormCudaExactBranchScales) { - TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); - TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); - TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); - TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); - TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); - TestStandardRMSBranchScales(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormWideSingleRowDeterministic) { - if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for deterministic RMS"; - constexpr int width = 8192; - constexpr float epsilon = 1e-6f; - std::vector input(width); - double square_sum = 0.0; - for (int i = 0; i < width; ++i) { - input[i] = std::sin(static_cast(i) * 0.71f) * (1 + i % 17); - const float square = input[i] * input[i]; - square_sum += square; - } - const float rms = std::sqrt(static_cast(square_sum / width) + epsilon); - std::vector expected(width); - for (int i = 0; i < width; ++i) expected[i] = input[i] / rms; - std::vector first; - for (int run = 0; run < 10; ++run) { - OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); - tester.AddAttribute("use_onnx_arithmetic", 1); - tester.AddAttribute("epsilon", epsilon); - tester.AddInput("X", {1, 1, width}, input); - tester.AddOptionalInputEdge(); - tester.AddOutput("Y", {1, 1, width}, expected, false, 1e-5f, 1e-5f); - tester.AddOutput("Y_fp32", {1, 1, width}, expected, false, 1e-5f, 1e-5f); - SessionOptions options; - options.use_deterministic_compute = true; - options.graph_optimization_level = TransformerLevel::Default; - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); - const auto& fetches = tester.GetFetches(); - ASSERT_EQ(fetches.size(), 2u); - const float* actual = fetches[0].Get().Data(); - if (first.empty()) { - first.assign(actual, actual + width); - } else { - EXPECT_EQ(std::memcmp(first.data(), actual, width * sizeof(float)), 0); - } - EXPECT_EQ(std::memcmp(actual, fetches[1].Get().Data(), width * sizeof(float)), 0); - } -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedRejectsInvalidSlice) { - if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for packed RMS"; - for (const auto& [start, width, mode, branches, message] : - std::vector>{ - {-1, 2, 1, 0, "nonnegative"}, - {0, -1, 1, 0, "nonnegative"}, - {1, 0, 1, 0, "nonnegative"}, - {0, 5, 1, 0, "exceeds"}, - {3, 2, 1, 0, "exceeds"}, - {0, 2, 0, 0, "requires"}, - {0, 2, 1, 2, "requires"}}) { - OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); - tester.AddAttribute("input_start", start); - tester.AddAttribute("input_width", width); - tester.AddAttribute("use_onnx_arithmetic", mode); - tester.AddAttribute("num_branches", branches); - tester.AddInput("X", {1, 1, 4}, {1, 2, 3, 4}); - tester.AddOptionalInputEdge(); - tester.AddOutput("Y", {1, 1, 2}, {0, 0}); - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectFailure, message, {}, nullptr, &providers); - } -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedEmptyWithoutScale) { - if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for packed RMS"; - OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); - tester.AddAttribute("use_onnx_arithmetic", 1); - tester.AddAttribute("input_start", 3); - tester.AddAttribute("input_width", 2); - tester.AddInput("X", {0, 3, 8}, {}); - tester.AddOptionalInputEdge(); - tester.AddOutput("Y", {0, 3, 2}, {}); - tester.AddOutput("Y_fp32", {0, 3, 2}, {}); - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedWithoutScale) { - if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for packed RMS"; - OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); - tester.AddAttribute("use_onnx_arithmetic", 1); - tester.AddAttribute("input_start", 1); - tester.AddAttribute("input_width", 2); - tester.AddAttribute("epsilon", 0.0f); - tester.AddInput("X", {1, 2, 4}, {99, 3, 4, 99, 99, 0, 2, 99}); - tester.AddOptionalInputEdge(); - tester.AddOutput("Y", {1, 2, 2}, - {3.0f / std::sqrt(12.5f), 4.0f / std::sqrt(12.5f), 0, 2.0f / std::sqrt(2.0f)}); - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedDoesNotFallBackToCpu) { - OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); - tester.AddAttribute("use_onnx_arithmetic", 1); - tester.AddAttribute("input_width", 2); - tester.AddInput("X", {1, 1, 4}, {1, 2, 3, 4}); - tester.AddOptionalInputEdge(); - tester.AddOutput("Y", {1, 1, 2}, {0, 0}); - std::vector> providers; - providers.push_back(DefaultCpuExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectFailure, "Could not find an implementation for BranchwiseRMSNorm(2)", - {}, nullptr, &providers); -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardCpuBranchScale) { - OpTester tester("BranchwiseRMSNorm", 1, kMSDomain); - tester.AddAttribute("use_onnx_arithmetic", 1); - tester.AddAttribute("num_branches", 2); - tester.AddAttribute("epsilon", 0.0f); - tester.AddInput("X", {1, 4}, {3.0f, 4.0f, 0.0f, 2.0f}); - tester.AddInput("scale", {2, 2}, {2.0f, 0.5f, 0.25f, 3.0f}); - tester.AddOutput("Y", {1, 4}, - {3.0f / std::sqrt(12.5f) * 2.0f, 4.0f / std::sqrt(12.5f) * 0.5f, - 0.0f, 2.0f / std::sqrt(2.0f) * 3.0f}); - std::vector> providers; - providers.push_back(DefaultCpuExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - -TEST(HyperConnectionOpsTest, BranchwiseRMSNormRejectsInvalidArithmeticMode) { - for (auto& provider : AvailableProviders()) { - if (provider->Type() == kWebGpuExecutionProvider) continue; - OpTester tester("BranchwiseRMSNorm", 1, kMSDomain); - tester.AddAttribute("use_onnx_arithmetic", 2); - tester.AddInput("X", {1, 2}, {1.0f, 2.0f}); - tester.AddOptionalInputEdge(); - tester.AddOutput("Y", {1, 2}, {0.0f, 0.0f}); - std::vector> providers; - providers.push_back(std::move(provider)); - tester.Run(OpTester::ExpectResult::kExpectFailure, "use_onnx_arithmetic must be 0 or 1", - {}, nullptr, &providers); - } -} - TEST(HyperConnectionOpsTest, ScaledSiLU) { RunOnAvailableProviders("ScaledSiLU", [&](OpTester& tester) { tester.AddAttribute("alpha", 0.5f); diff --git a/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc b/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc index 73ce50f748149..73c32a8ac4878 100644 --- a/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc +++ b/onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc @@ -84,40 +84,63 @@ TEST(RMSNormRotaryTest, RMSNormRotaryExactUnfusedParity) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); }; if (CudaProvider().empty()) GTEST_SKIP(); - OpTester original("BranchwiseRMSNorm", 2, kMSDomain); - original.AddAttribute("use_onnx_arithmetic", 1); - original.AddAttribute("input_width", width); - original.AddAttribute("input_start", start); - original.AddAttribute("epsilon", 1e-6f); - original.AddInput("X", input_shape, input); - original.AddInput("scale", {width}, scale); - original.AddInput("mode", {1}, {dense}); - original.AddOutput("Y", norm_shape, norm); - original.AddOutput("Y_fp32", norm_shape, normalized); - original.SetCustomOutputVerifier([&](const std::vector& values, const std::string&) { - if (norm.empty()) return; - const auto& y = values[0].Get(); - norm.assign(y.Data(), y.Data() + norm.size()); - const auto& widened = values[1].Get(); - normalized.assign(widened.Data(), widened.Data() + normalized.size()); - }); - run(original); - OpTester rope("RotaryEmbedding", 2, kMSDomain); - rope.AddAttribute("num_heads", heads); - rope.AddAttribute("rotary_embedding_dim", rotary); - rope.AddAttribute("use_fp32_arithmetic", fp32 ? 1 : 0); - rope.AddInput("X", output_shape, norm); - rope.AddInput("position", position_shape, positions); - rope.AddInput("cos", {cache_rows, rotary / 2}, cosine); - rope.AddInput("sin", {cache_rows, rotary / 2}, sine); - rope.AddInput("mode", {1}, {dense}); - rope.AddOutput("Y", output_shape, expected); - rope.SetCustomOutputVerifier([&](const std::vector& values, const std::string&) { - if (expected.empty()) return; - const auto& y = values[0].Get(); - expected.assign(y.Data(), y.Data() + expected.size()); - }); - if (sequence > 0) run(rope); + 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); diff --git a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc index 14bc41c6da47b..6696790410603 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc +++ b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc @@ -1237,9 +1237,7 @@ static void BuildRotaryEmbeddingGQAFusionGraph(ModelTestBuilder& builder, int64_t q_interleaved = 0, int64_t k_interleaved = 0, int64_t rotary_embedding_dim = 0, - int64_t gqa_num_heads = 2, - bool q_fp32_arithmetic = false, - bool k_fp32_arithmetic = false) { + int64_t gqa_num_heads = 2) { constexpr int64_t batch_size = 1; constexpr int64_t sequence_length = 2; constexpr int64_t input_hidden_size = 8; @@ -1309,10 +1307,6 @@ static void BuildRotaryEmbeddingGQAFusionGraph(ModelTestBuilder& builder, k_rotary.AddAttribute("num_heads", kv_num_heads); k_rotary.AddAttribute("interleaved", k_interleaved); k_rotary.AddAttribute("rotary_embedding_dim", rotary_embedding_dim); - if (q_fp32_arithmetic || k_fp32_arithmetic) { - q_rotary.AddAttribute("use_fp32_arithmetic", static_cast(q_fp32_arithmetic)); - k_rotary.AddAttribute("use_fp32_arithmetic", static_cast(k_fp32_arithmetic)); - } NodeArg* past_key = builder.MakeInput({{batch_size, kv_num_heads, max_sequence_length, head_size}}); @@ -1724,28 +1718,6 @@ TEST_F(GraphTransformationTests, GroupQueryAttentionFusionMsRotaryEmbeddingForwa check_fused_graph)); } -TEST_F(GraphTransformationTests, GroupQueryAttentionFusionPreservesExactRotaryArithmetic) { - for (const auto& [q_exact, k_exact] : - {std::pair{false, false}, std::pair{true, false}, std::pair{false, true}, std::pair{true, true}}) { - Model model("ExactRotaryGQAFusion", false, ModelMetaData(), PathString(), - IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 23}, {kMSDomain, 2}}, {}, *logger_); - Graph& graph = model.MainGraph(); - ModelTestBuilder builder(graph); - BuildRotaryEmbeddingGQAFusionGraph(builder, RotaryEmbeddingDomain::kMS, true, 1, 1, 0, 2, - q_exact, k_exact); - builder.SetGraphOutputs(); - ASSERT_STATUS_OK(graph.Resolve()); - GraphTransformerManager transformers{3}; - ASSERT_STATUS_OK(transformers.Register(std::make_unique(), TransformerLevel::Level2)); - ASSERT_STATUS_OK(transformers.ApplyTransformers(graph, TransformerLevel::Level2, *logger_)); - if (q_exact || k_exact) { - ASSERT_STATUS_OK(CheckMsRotaryEmbeddingGQANotFused(graph)); - } else { - ASSERT_STATUS_OK(CheckMsRotaryEmbeddingGQAFused(graph, 1)); - } - } -} - TEST_F(GraphTransformationTests, GroupQueryAttentionFusionSkipsMismatchedProjectionSizesTest) { auto build_test_case = [](ModelTestBuilder& builder) { BuildRotaryEmbeddingGQAFusionGraph(builder, RotaryEmbeddingDomain::kMS, true, 0, 0, 0, 3); diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 207038805f624..92f803030ada4 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -5082,7 +5082,7 @@ void RegisterOrtOpSchemas() { auto& domainToVersionRangeInstance = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance(); if (domainToVersionRangeInstance.Map().find(onnxruntime::kMSDomain) == domainToVersionRangeInstance.Map().end()) { // External shared providers may have already added kMSDomain - domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 2); + domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 1); } domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSExperimentalDomain, 1, 1); domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSNchwcDomain, 1, 1); From 6c2f985b62926e4892a2c1a54b07547137ad0326 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 04:09:32 +0000 Subject: [PATCH 8/8] fix(docs): sync RMSNormRotary schema docs in ContribOperators.md Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index db014df0124b8..6dcb758a234fe 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -7110,17 +7110,17 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
epsilon : float (default is 1e-06)
+
epsilon : float
Finite nonnegative RMS epsilon.
-
flatten_batch : int (default is 0)
+
flatten_batch : int
One flattens B*S to sequence and outputs [1,B*S,heads*width].
-
input_start : int (default is 0)
+
input_start : int
Nonnegative start of normalized last-axis slice.
-
input_width : int (default is 0)
+
input_width : int
Normalized width; zero uses the full last dimension.
-
rotary_embedding_dim : int
+
rotary_embedding_dim : int (required)
Positive even rotary width, no larger than normalized width.
-
use_fp32_arithmetic : int (default is 1)
+
use_fp32_arithmetic : int
Decode rotary FP32 (1) or BF16 intermediates (0).