From ceb91dd2b88f54804930cd9a1efac3c646c3c093 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 19:12:48 +0000 Subject: [PATCH 01/19] Add fused hyper-connection and tensor CUDA operators --- docs/ContribOperators.md | 87 ++ .../contrib_ops/cpu/bert/engram_gate.cc | 25 +- .../contrib_ops/cuda/bert/engram_gate.cc | 27 +- .../contrib_ops/cuda/bert/engram_gate_impl.cu | 7 +- .../contrib_ops/cuda/bert/engram_gate_impl.h | 3 +- .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 260 +++++ .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 352 +++++++ .../contrib_ops/cuda/bert/fused_tensor_ops.h | 97 ++ .../bert/hyper_connection_low_rank_mix.cc | 165 ++++ .../cuda/bert/hyper_connection_low_rank_mix.h | 25 + .../hyper_connection_low_rank_mix_impl.cu | 888 ++++++++++++++++++ .../bert/hyper_connection_low_rank_mix_impl.h | 67 ++ .../contrib_ops/cuda/cuda_contrib_kernels.cc | 20 + .../core/graph/contrib_ops/bert_defs.cc | 255 ++++- onnxruntime/core/graph/contrib_ops/ms_opset.h | 12 + .../test/contrib_ops/engram_ops_test.cc | 53 +- .../test/contrib_ops/fused_tensor_ops_test.cc | 283 ++++++ .../hyper_connection_low_rank_mix_op_test.cc | 285 ++++++ .../contrib_ops/hyper_connection_op_test.cc | 252 +++++ .../test_hyper_connection_low_rank_mix.py | 311 ++++++ 20 files changed, 3431 insertions(+), 43 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc create mode 100644 onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu create mode 100644 onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h create mode 100644 onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc create mode 100644 onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h create mode 100644 onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu create mode 100644 onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h create mode 100644 onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc create mode 100644 onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc create mode 100644 onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 2703de7cc1fae..1015b41e18389 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -3,6 +3,7 @@ Do not modify directly.* * com.microsoft + * com.microsoft.HyperConnectionLowRankMix * com.microsoft.Attention * com.microsoft.AttnLSTM * com.microsoft.BeamSearch @@ -8772,3 +8773,89 @@ No versioning maintained for experimental ops.
Constrain input and output types to float32 tensors.
+### **com.microsoft.HyperConnectionLowRankMix** + + One hyper-connection block with a low-rank gate network, fused into one operator. C = num_branches streams of + width H are kept as one flattened (..., C * H) tensor. With the optional prologue the operator first applies the + post-mix of the previous block (identity stream mix): + + S[..., c, h] = streams[..., c, h] + prev_post_gate[..., c] * prev_block_output[..., h] + + and otherwise S = streams. It then computes, with `rms` taken per branch over H: + + hn[..., c, h] = S[..., c, h] / sqrt(mean_h(S[..., c, :]^2) + epsilon) * norm_scale[c, h] + d = ScaledSiLU(hn @ down_weight, silu_alpha) (..., R) + g[..., c, h] = sigmoid(d @ up_weight) per stream and feature + pre_output = reduction_scale * sum_c(hn[..., c, h] * g[..., c, h]) (..., H) + post_gate = sigmoid((hn @ inject_weight) / gate_div) * gate_mul (..., C) + + pre_output is the input of the next transformer block and post_gate is prev_post_gate of the next call. + streams_out is S (the streams after the prologue; a copy of streams when there is no prologue). + + Numerics contract. Matrix products and reductions accumulate in float32. The result of every step is rounded + to T once, as the output of the equivalent standard operator would be: + S = cast_T(streams + prev_post_gate * prev_block_output), hn = cast_T(S * inv_rms * norm_scale), + down = cast_T(hn @ down_weight), Z = cast_T(silu_alpha * down), SG = cast_T(sigmoid(Z)), hidden = cast_T(Z * SG), + up = cast_T(hidden @ up_weight), g = cast_T(sigmoid(up)), pre_output = cast_T(reduction_scale * sum_c(hn * g)), + inject = cast_T(hn @ inject_weight), post_gate = cast_T(cast_T(sigmoid(cast_T(inject / gate_div))) * gate_mul). + sigmoid is evaluated in float32 with expf and rounded once to T; it does not reproduce a Sigmoid kernel that + rounds to T after every elementary operation. The accumulation order of the matrix products is fixed, so the + results are deterministic, but they are not bitwise identical to cuBLAS. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
epsilon : float
+
Epsilon added to the mean square of each branch before the reciprocal square root.
+
gate_div : float
+
Divisor applied to the inject logits before the sigmoid. Must not be zero.
+
gate_mul : float
+
Multiplier applied to the sigmoid of the inject logits.
+
num_branches : int (required)
+
Number of residual streams C. The last dimension of streams is C * H.
+
reduction_scale : float
+
Multiplier applied to the sum over branches of pre_output.
+
silu_alpha : float
+
Scale of the ScaledSiLU between the down and up projections.
+
+ +#### Inputs (5 - 7) + +
+
streams : T
+
Flattened streams (..., C * H). With the prologue these are the streams before the previous block's post-mix.
+
norm_scale : T
+
Per-branch RMS norm weight with shape (C * H).
+
down_weight : T
+
Down projection with shape (C * H, R), row-major [K, N].
+
up_weight : T
+
Up projection with shape (R, C * H), row-major [K, N].
+
inject_weight : T
+
Post-gate projection with shape (C * H, C), row-major [K, N].
+
prev_block_output (optional) : T
+
Optional output of the previous block with shape (..., H). Provide together with prev_post_gate.
+
prev_post_gate (optional) : T
+
Optional post-mix gate of the previous block with shape (..., C). Provide together with prev_block_output.
+
+ +#### Outputs (2 - 3) + +
+
pre_output : T
+
Block input with shape (..., H).
+
post_gate : T
+
Post-mix gate for the next call with shape (..., C).
+
streams_out (optional) : T
+
Optional streams after the prologue, with the same shape as streams. Equal to streams when the prologue inputs are absent.
+
+ +#### Type Constraints + +
+
T : tensor(float), tensor(float16), tensor(bfloat16)
+
Constrain all floating-point inputs and outputs to the same type.
+
diff --git a/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc b/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc index c86c6005533cb..edae2a9685d9f 100644 --- a/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc +++ b/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc @@ -45,16 +45,23 @@ Status EngramGate::Compute(OpKernelContext* context) const { const Tensor* conv_norm_scale = context->Input(5); const TensorShape& key_shape = key->Shape(); - ORT_RETURN_IF_NOT(key_shape.NumDimensions() == 4, - "key must have shape (batch_size, sequence_length, hc_mult, hidden_size)"); - const int64_t batch_size = key_shape[0]; - const int64_t sequence_length = key_shape[1]; - const int64_t hc_mult = key_shape[2]; - const int64_t hidden_size = key_shape[3]; + const size_t key_rank = key_shape.NumDimensions(); + ORT_RETURN_IF_NOT(key_rank == 3 || key_rank == 4, + "key must have shape (total_tokens, hc_mult, hidden_size) or " + "(batch_size, sequence_length, hc_mult, hidden_size)"); + const size_t token_dims = key_rank - 2; + const int64_t hc_mult = key_shape[token_dims]; + const int64_t hidden_size = key_shape[token_dims + 1]; + TensorShapeVector expected_value_dims; + expected_value_dims.reserve(key_rank - 1); + for (size_t i = 0; i < token_dims; ++i) { + expected_value_dims.push_back(key_shape[i]); + } + expected_value_dims.push_back(hidden_size); ORT_RETURN_IF_NOT(query->Shape() == key_shape, "query must have the same shape as key"); - ORT_RETURN_IF_NOT(value->Shape() == TensorShape({batch_size, sequence_length, hidden_size}), - "value must have shape (batch_size, sequence_length, hidden_size)"); + ORT_RETURN_IF_NOT(value->Shape() == TensorShape(expected_value_dims), + "value must match key's token dimensions and hidden_size"); ORT_RETURN_IF_NOT(key_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}), "key_norm_scale must have shape (hc_mult, hidden_size)"); ORT_RETURN_IF_NOT(query_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}), @@ -81,7 +88,7 @@ Status EngramGate::Compute(OpKernelContext* context) const { T* output_data = output->MutableData(); T* output_normed_data = output_normed == nullptr ? nullptr : output_normed->MutableData(); - const int64_t rows = batch_size * sequence_length * hc_mult; + const int64_t rows = key_shape.SizeToDimension(key_rank - 1); ThreadPool::TryParallelFor( context->GetOperatorThreadPool(), narrow(rows), static_cast(2 * hidden_size + 32), diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc b/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc index 825a35aba4121..e870950062761 100644 --- a/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc +++ b/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc @@ -45,16 +45,24 @@ Status EngramGate::ComputeInternal(OpKernelContext* context) const { const Tensor* conv_norm_scale = context->Input(5); const TensorShape& key_shape = key->Shape(); - ORT_RETURN_IF_NOT(key_shape.NumDimensions() == 4, - "key must have shape (batch_size, sequence_length, hc_mult, hidden_size)"); - const int64_t batch_size = key_shape[0]; - const int64_t sequence_length = key_shape[1]; - const int64_t hc_mult = key_shape[2]; - const int64_t hidden_size = key_shape[3]; + const size_t key_rank = key_shape.NumDimensions(); + ORT_RETURN_IF_NOT(key_rank == 3 || key_rank == 4, + "key must have shape (total_tokens, hc_mult, hidden_size) or " + "(batch_size, sequence_length, hc_mult, hidden_size)"); + const size_t token_dims = key_rank - 2; + const int64_t token_count = key_shape.SizeToDimension(token_dims); + const int64_t hc_mult = key_shape[token_dims]; + const int64_t hidden_size = key_shape[token_dims + 1]; + TensorShapeVector expected_value_dims; + expected_value_dims.reserve(key_rank - 1); + for (size_t i = 0; i < token_dims; ++i) { + expected_value_dims.push_back(key_shape[i]); + } + expected_value_dims.push_back(hidden_size); ORT_RETURN_IF_NOT(query->Shape() == key_shape, "query must have the same shape as key"); - ORT_RETURN_IF_NOT(value->Shape() == TensorShape({batch_size, sequence_length, hidden_size}), - "value must have shape (batch_size, sequence_length, hidden_size)"); + ORT_RETURN_IF_NOT(value->Shape() == TensorShape(expected_value_dims), + "value must match key's token dimensions and hidden_size"); ORT_RETURN_IF_NOT(key_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}), "key_norm_scale must have shape (hc_mult, hidden_size)"); ORT_RETURN_IF_NOT(query_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}), @@ -82,8 +90,7 @@ Status EngramGate::ComputeInternal(OpKernelContext* context) const { conv_norm_scale == nullptr ? nullptr : reinterpret_cast(conv_norm_scale->Data()), reinterpret_cast(output->MutableData()), output_normed == nullptr ? nullptr : reinterpret_cast(output_normed->MutableData()), - batch_size, - sequence_length, + token_count, hc_mult, hidden_size, epsilon_); diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu index 788d9bf9ed72e..fbcb761667430 100644 --- a/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu @@ -96,12 +96,11 @@ Status LaunchEngramGateKernel( const T* conv_norm_scale, T* output, T* output_normed, - int64_t batch_size, - int64_t sequence_length, + int64_t token_count, int64_t hc_mult, int64_t hidden_size, float epsilon) { - const int64_t rows = batch_size * sequence_length * hc_mult; + const int64_t rows = token_count * hc_mult; if (rows == 0 || hidden_size == 0) { return Status::OK(); } @@ -115,7 +114,7 @@ Status LaunchEngramGateKernel( #define INSTANTIATE_ENGRAM_GATE(T) \ template Status LaunchEngramGateKernel(cudaStream_t, const T*, const T*, const T*, const T*, const T*, \ - const T*, T*, T*, int64_t, int64_t, int64_t, int64_t, float); + const T*, T*, T*, int64_t, int64_t, int64_t, float); INSTANTIATE_ENGRAM_GATE(float) INSTANTIATE_ENGRAM_GATE(half) diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h index fdacbbd393600..707a49daeb221 100644 --- a/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h @@ -21,8 +21,7 @@ Status LaunchEngramGateKernel( const T* conv_norm_scale, T* output, T* output_normed, - int64_t batch_size, - int64_t sequence_length, + int64_t token_count, int64_t hc_mult, int64_t hidden_size, float epsilon); diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc new file mode 100644 index 0000000000000..555c98f1fb16d --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -0,0 +1,260 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/fused_tensor_ops.h" + +#include +#include +#include + +#include "core/common/float8.h" +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/cuda_type_conversion.h" + +namespace onnxruntime::contrib::cuda { + +#define REGISTER_ARGMAX_LAST_AXIS(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + ArgMaxLastAxis, kMSDomain, 1, T, kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ + ArgMaxLastAxis); + +REGISTER_ARGMAX_LAST_AXIS(float) +REGISTER_ARGMAX_LAST_AXIS(MLFloat16) +REGISTER_ARGMAX_LAST_AXIS(BFloat16) + +#undef REGISTER_ARGMAX_LAST_AXIS + +ONNX_OPERATOR_KERNEL_EX( + EmbeddingRMSNorm, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("F", DataTypeImpl::GetTensorType()), + EmbeddingRMSNorm); + +ONNX_OPERATOR_KERNEL_EX( + RowSelect, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + RowSelect); + +ONNX_OPERATOR_KERNEL_EX( + SigmoidGateMul, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + SigmoidGateMul); + +#if !defined(DISABLE_FLOAT8_TYPES) +ONNX_OPERATOR_TYPED_KERNEL_EX( + DequantizeLinear, kMSDomain, 1, Float8E4M3FN_BFloat16, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("TS", DataTypeImpl::GetTensorType()), + DequantizeLinearFp8Bf16); +#endif + +ONNX_OPERATOR_KERNEL_EX( + UntileHyperConnectionWeights, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + UntileHyperConnectionWeights); + +template +ArgMaxLastAxis::ArgMaxLastAxis(const OpKernelInfo& info) + : CudaKernel(info), keepdims_(info.GetAttrOrDefault("keepdims", 0) != 0) {} + +template +Status ArgMaxLastAxis::ComputeInternal(OpKernelContext* context) const { + const Tensor* input = context->Input(0); + const auto& shape = input->Shape(); + ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1 && + shape[shape.NumDimensions() - 1] > 0, + "ArgMaxLastAxis requires a non-empty final dimension"); + const int64_t width = shape[shape.NumDimensions() - 1]; + const int64_t rows = shape.Size() / width; + TensorShapeVector output_dims(shape.GetDims().begin(), shape.GetDims().end() - 1); + if (keepdims_) output_dims.push_back(1); + Tensor* output = context->Output(0, TensorShape(output_dims)); + const size_t partial_count = SafeInt(rows) * GetArgMaxLastAxisPartitions(width); + auto partial_values = GetScratchBuffer(partial_count, GetComputeStream(context)); + auto partial_indices = GetScratchBuffer(partial_count, GetComputeStream(context)); + using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; + return LaunchArgMaxLastAxis(Stream(context), reinterpret_cast(input->Data()), + input->GetElementType(), rows, width, + output->MutableData(), partial_values.get(), partial_indices.get()); +} + +EmbeddingRMSNorm::EmbeddingRMSNorm(const OpKernelInfo& info) + : CudaKernel(info), epsilon_(info.GetAttrOrDefault("epsilon", 1e-6f)) { + ORT_ENFORCE(epsilon_ >= 0.0f, "EmbeddingRMSNorm epsilon must be non-negative"); +} + +Status EmbeddingRMSNorm::ComputeInternal(OpKernelContext* context) const { + const Tensor* ids = context->Input(0); + const Tensor* table = context->Input(1); + const Tensor* weight = context->Input(2); + ORT_RETURN_IF_NOT(table->Shape().NumDimensions() == 2, + "EmbeddingRMSNorm table must have shape [vocabulary, hidden]"); + const int64_t hidden_size = table->Shape()[1]; + ORT_RETURN_IF_NOT(weight->Shape().NumDimensions() == 1 && + weight->Shape()[0] == hidden_size, + "EmbeddingRMSNorm weight must match the embedding hidden size"); + TensorShapeVector output_dims(ids->Shape().GetDims().begin(), ids->Shape().GetDims().end()); + output_dims.push_back(hidden_size); + Tensor* output = context->Output(0, TensorShape(output_dims)); + return LaunchEmbeddingRMSNorm(Stream(context), ids->Data(), + reinterpret_cast(table->Data()), + weight->Data(), epsilon_, + reinterpret_cast(output->MutableData()), + ids->Shape().Size(), table->Shape()[0], hidden_size); +} + +Status RowSelect::ComputeInternal(OpKernelContext* context) const { + const Tensor* selection = context->Input(0); + ORT_RETURN_IF_NOT(selection->Shape().Size() == 1, "RowSelect selection must have one element"); + std::array inputs{}; + std::array outputs{}; + std::array row_sizes{}; + for (int i = 0; i < 3; ++i) { + const Tensor* input = context->Input(i + 1); + const auto& shape = input->Shape(); + ORT_RETURN_IF_NOT(shape.NumDimensions() >= 2 && shape[0] == 1 && shape[1] == 2, + "RowSelect inputs must have shape [1, 2, ...]"); + row_sizes[i] = shape.Size() / 2; + TensorShapeVector output_dims(shape.GetDims().begin(), shape.GetDims().end()); + output_dims[1] = 1; + Tensor* output = context->Output(i, TensorShape(output_dims)); + inputs[i] = reinterpret_cast(input->Data()); + outputs[i] = reinterpret_cast(output->MutableData()); + } + return LaunchRowSelect(Stream(context), selection->Data(), inputs.data(), + outputs.data(), row_sizes.data()); +} + +SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) + : CudaKernel(info), + head_dim_(static_cast(info.GetAttrOrDefault("head_dim", 0))), + heads_(static_cast(info.GetAttrOrDefault("num_heads", 0))), + q_stride_(static_cast(info.GetAttrOrDefault("q_stride", 0))), + tok_stride_(info.GetAttrOrDefault("tok_stride", 0)), + q_offset_(info.GetAttrOrDefault("q_offset", 0)), + broadcast_gate_(info.GetAttrOrDefault("broadcast_gate", 0) == 1), + sigmoid_in_fp32_(info.GetAttrOrDefault("sigmoid_in_fp32", 1) == 1), + silu_gate_(info.GetAttrOrDefault("silu_gate", 0) == 1) { + const int64_t broadcast = info.GetAttrOrDefault("broadcast_gate", 0); + ORT_ENFORCE(broadcast == 0 || broadcast == 1, "SigmoidGateMul broadcast_gate must be zero or one"); + const int64_t fp32 = info.GetAttrOrDefault("sigmoid_in_fp32", 1); + ORT_ENFORCE(fp32 == 0 || fp32 == 1, "SigmoidGateMul sigmoid_in_fp32 must be zero or one"); + const int64_t silu = info.GetAttrOrDefault("silu_gate", 0); + ORT_ENFORCE(silu == 0 || silu == 1, "SigmoidGateMul silu_gate must be zero or one"); + ORT_ENFORCE(!(broadcast_gate_ && silu_gate_), "SigmoidGateMul broadcast_gate and silu_gate are exclusive"); + ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && q_stride_ >= 2 * head_dim_ && tok_stride_ > 0 && q_offset_ >= 0), + "SigmoidGateMul received invalid layout attributes"); + ORT_ENFORCE(!(broadcast_gate_ || silu_gate_) || (head_dim_ == 0 && heads_ == 0 && q_stride_ == 0 && + tok_stride_ == 0 && q_offset_ == 0), + "SigmoidGateMul broadcast_gate/silu_gate cannot be combined with packed layout attributes"); +} + +Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { + const Tensor* attention = context->Input(0); + const Tensor* projection = context->Input(1); + const Tensor* residual = context->Input(2); + const int64_t count = attention->Shape().Size(); + ORT_RETURN_IF_NOT(residual == nullptr || residual->Shape() == attention->Shape(), + "SigmoidGateMul residual must have the same shape as attention"); + int64_t gate_row_width = 0; + if (silu_gate_) { + ORT_RETURN_IF_NOT(projection->Shape() == attention->Shape(), + "SigmoidGateMul silu gate must have the same shape as attention"); + } else if (broadcast_gate_) { + TensorShapeVector expected(attention->Shape().GetDims().begin(), attention->Shape().GetDims().end()); + ORT_RETURN_IF_NOT(!expected.empty() && expected.back() > 0, + "SigmoidGateMul broadcast_gate requires a positive final attention dimension"); + gate_row_width = expected.back(); + expected.back() = 1; + ORT_RETURN_IF_NOT(projection->Shape() == TensorShape(expected), + "SigmoidGateMul broadcast gate must match attention with a singleton final dimension"); + } else { + const int64_t row_width = static_cast(heads_) * head_dim_; + ORT_RETURN_IF_NOT(count % row_width == 0, + "SigmoidGateMul attention size must be divisible by heads * head_dim"); + const int64_t tokens = count / row_width; + ORT_RETURN_IF_NOT(projection->Shape().Size() >= + (tokens - 1) * tok_stride_ + q_offset_ + + static_cast(heads_) * q_stride_, + "SigmoidGateMul projection is too small for the configured layout"); + } + Tensor* output = context->Output(0, attention->Shape()); + return LaunchSigmoidGateMul( + Stream(context), reinterpret_cast(attention->Data()), + reinterpret_cast(projection->Data()), + reinterpret_cast(output->MutableData()), count, head_dim_, heads_, + q_stride_, tok_stride_, q_offset_, gate_row_width, + residual == nullptr ? nullptr : reinterpret_cast(residual->Data()), + sigmoid_in_fp32_, silu_gate_); +} + +DequantizeLinearFp8Bf16::DequantizeLinearFp8Bf16(const OpKernelInfo& info) + : CudaKernel(info), axis_(info.GetAttrOrDefault("axis", 1)) {} + +Status DequantizeLinearFp8Bf16::ComputeInternal(OpKernelContext* context) const { + const Tensor* input = context->Input(0); + const Tensor* scales = context->Input(1); + const Tensor* zero_point = context->Input(2); + const auto& input_shape = input->Shape(); + ORT_RETURN_IF_NOT(input_shape.NumDimensions() > 0, + "DequantizeLinear FP8/BF16 input must have rank at least one"); + int64_t axis = axis_ < 0 ? axis_ + static_cast(input_shape.NumDimensions()) : axis_; + + int64_t axis_size = 1; + int64_t inner_size = 1; + if (scales->Shape().Size() != 1) { + ORT_RETURN_IF_NOT(axis >= 0 && axis < static_cast(input_shape.NumDimensions()), + "DequantizeLinear axis is out of range"); + ORT_RETURN_IF_NOT(scales->Shape().NumDimensions() == 1 && + scales->Shape().Size() == input_shape[axis], + "DequantizeLinear scale must be scalar or match the selected axis"); + axis_size = scales->Shape().Size(); + inner_size = input_shape.SizeFromDimension(static_cast(axis + 1)); + } + if (zero_point != nullptr) { + ORT_RETURN_IF_NOT(zero_point->Shape().Size() == 1 || + zero_point->Shape() == scales->Shape(), + "DequantizeLinear zero point must be scalar or match the scale"); + } + + Tensor* output = context->Output(0, input_shape); + return LaunchDequantizeLinearFp8Bf16( + Stream(context), reinterpret_cast(input->DataRaw()), + zero_point == nullptr ? nullptr : reinterpret_cast(zero_point->DataRaw()), + scales->Data(), + reinterpret_cast(output->MutableData()), input_shape.Size(), + axis_size, inner_size, zero_point == nullptr || zero_point->Shape().Size() == 1); +} + +UntileHyperConnectionWeights::UntileHyperConnectionWeights(const OpKernelInfo& info) + : CudaKernel(info), kind_(static_cast(info.GetAttrOrDefault("kind", -1))) { + ORT_ENFORCE(kind_ == 0 || kind_ == 1, "UntileHyperConnectionWeights kind must be 0 or 1"); +} + +Status UntileHyperConnectionWeights::ComputeInternal(OpKernelContext* context) const { + const Tensor* input = context->Input(0); + ORT_RETURN_IF_NOT(input->Shape() == TensorShape({160, 20480}), + "UntileHyperConnectionWeights input must have shape [160, 20480]"); + const TensorShape output_shape = kind_ == 0 ? TensorShape({10240, 320}) : TensorShape({320, 10240}); + Tensor* output = context->Output(0, output_shape); + return LaunchUntileHyperConnectionWeights( + Stream(context), reinterpret_cast(input->Data()), + reinterpret_cast(output->MutableData()), kind_); +} + +template class ArgMaxLastAxis; +template class ArgMaxLastAxis; +template class ArgMaxLastAxis; + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu new file mode 100644 index 0000000000000..789edc4a26615 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -0,0 +1,352 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/fused_tensor_ops.h" + +#include +#include +#include +#include +#include + +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/cu_inc/common.cuh" + +namespace onnxruntime::contrib::cuda { +namespace { + +constexpr int kThreads = 256; +constexpr int kArgmaxThreads = 256; + +__device__ __forceinline__ float ToFloat(float value) { return value; } +__device__ __forceinline__ float ToFloat(half value) { return __half2float(value); } +__device__ __forceinline__ float ToFloat(__nv_bfloat16 value) { return __bfloat162float(value); } + +template +__global__ void ArgmaxRowsKernel(const T* input, const int64_t* input_indices, + int64_t width, int64_t partitions, float* partial_values, + int64_t* output) { + __shared__ float values[kArgmaxThreads]; + __shared__ int64_t indices[kArgmaxThreads]; + const int tid = threadIdx.x; + const int64_t row = blockIdx.x / partitions; + const int64_t partition = blockIdx.x % partitions; + const int64_t end = partitions == 1 ? width : min(width, (partition + 1) * 4096); + float best_value = -CUDART_INF_F; + int64_t best_index = INT64_MAX; + for (int64_t col = partition * 4096 + tid; col < end; col += blockDim.x) { + const float value = ToFloat(input[row * width + col]); + const int64_t index = input_indices == nullptr ? col : input_indices[row * width + col]; + if (value > best_value || (value == best_value && index < best_index)) { + best_value = value; + best_index = index; + } + } + values[tid] = best_value; + indices[tid] = best_index; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + const float other_value = values[tid + stride]; + const int64_t other_index = indices[tid + stride]; + if (other_value > values[tid] || + (other_value == values[tid] && other_index < indices[tid])) { + values[tid] = other_value; + indices[tid] = other_index; + } + } + __syncthreads(); + } + if (tid == 0) { + output[blockIdx.x] = indices[0]; + if (partial_values != nullptr) partial_values[blockIdx.x] = values[0]; + } +} + +template +Status LaunchTypedArgmax(cudaStream_t stream, const T* input, int64_t rows, + int64_t width, int64_t* output, + float* partial_values, int64_t* partial_indices) { + if (rows == 0) return Status::OK(); + const int64_t partitions = GetArgMaxLastAxisPartitions(width); + ORT_RETURN_IF_NOT(rows <= std::numeric_limits::max() / partitions, + "ArgMaxLastAxis has too many rows for a CUDA grid"); + ArgmaxRowsKernel<<(rows * partitions), kArgmaxThreads, 0, stream>>>( + input, nullptr, width, partitions, partitions == 1 ? nullptr : partial_values, + partitions == 1 ? output : partial_indices); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + if (partitions > 1) { + ArgmaxRowsKernel<<(rows), kArgmaxThreads, 0, stream>>>( + partial_values, partial_indices, partitions, 1, nullptr, output); + } + return CUDA_CALL(cudaGetLastError()); +} + +__global__ void EmbedNormKernel(const int64_t* ids, const uint16_t* table, + const float* weight, float epsilon, + uint16_t* output, int64_t rows, int64_t hidden_size) { + __shared__ float sums[512]; + const int64_t row_index = blockIdx.x; + const int64_t token_id = ids[row_index]; + const uint16_t* embedding = table + token_id * hidden_size; + float sum = 0.0f; + for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { + const float value = __uint_as_float(static_cast(embedding[col]) << 16); + sum += value * value; + } + sums[threadIdx.x] = sum; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + sums[threadIdx.x] += sums[threadIdx.x + stride]; + } + __syncthreads(); + } + const float scale = 1.0f / sqrtf(sums[0] / static_cast(hidden_size) + epsilon); + for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { + const float value = __uint_as_float(static_cast(embedding[col]) << 16); + const float normalized = (value * scale) * weight[col]; + output[row_index * hidden_size + col] = + __bfloat16_as_ushort(__float2bfloat16_rn(normalized)); + } + (void)rows; +} + +__global__ void RowSelectKernel(const int64_t* selection, + const uint16_t* input0, const uint16_t* input1, + const uint16_t* input2, uint16_t* output0, + uint16_t* output1, uint16_t* output2, + int64_t row_size0, int64_t row_size1, + int64_t row_size2) { + const int which = blockIdx.y; + const uint16_t* input = which == 0 ? input0 : (which == 1 ? input1 : input2); + uint16_t* output = which == 0 ? output0 : (which == 1 ? output1 : output2); + const int64_t row_size = which == 0 ? row_size0 : (which == 1 ? row_size1 : row_size2); + const int64_t selected_row = selection[0] != 0 ? 1 : 0; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + i < row_size; i += static_cast(gridDim.x) * blockDim.x) { + output[i] = input[selected_row * row_size + i]; + } +} + +__device__ __forceinline__ float StableSigmoid(float value) { + if (value >= 0.0f) { + return 1.0f / (1.0f + expf(-value)); + } + const float exp_value = expf(value); + return exp_value / (1.0f + exp_value); +} + +template +__global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention, + const __nv_bfloat16* projection, + __nv_bfloat16* output, int64_t count, + int head_dim, int heads, int q_stride, + int64_t token_stride, int64_t q_offset, + int64_t gate_row_width, const __nv_bfloat16* residual, + bool sigmoid_in_fp32) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count) return; + int64_t gate_index; + if constexpr (SiluGate) { + gate_index = i; + } else if constexpr (BroadcastGate) { + gate_index = i / gate_row_width; + } else { + const int64_t row = i / head_dim; + const int64_t token = row / heads; + const int head = static_cast(row % heads); + const int dim = static_cast(i % head_dim); + gate_index = token * token_stride + q_offset + + static_cast(head) * q_stride + head_dim + dim; + } + const float gate = __bfloat162float(projection[gate_index]); + __nv_bfloat16 sigmoid_bf16; + if (sigmoid_in_fp32) { + sigmoid_bf16 = __float2bfloat16_rn(StableSigmoid(gate)); + } else { + using namespace onnxruntime::cuda; + const BFloat16 a(gate); + // Match CUDA Sigmoid's input-type intermediates, including its negative-input subtraction. + const BFloat16 value = a > BFloat16(0.0f) + ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) + : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a)); + sigmoid_bf16 = __float2bfloat16_rn(static_cast(value)); + } + if constexpr (SiluGate) { + sigmoid_bf16 = __float2bfloat16_rn(gate * __bfloat162float(sigmoid_bf16)); + } + __nv_bfloat16 value = __float2bfloat16_rn( + __bfloat162float(attention[i]) * __bfloat162float(sigmoid_bf16)); + if constexpr (HasResidual) { + value = __float2bfloat16_rn(__bfloat162float(value) + __bfloat162float(residual[i])); + } + output[i] = value; +} + +__device__ __forceinline__ float DecodeE4M3FN(uint8_t value) { + const float sign = (value & 0x80u) != 0 ? -1.0f : 1.0f; + const int exponent = (value >> 3) & 0x0f; + const int mantissa = value & 0x07; + if (exponent == 0) { + return sign * ldexpf(static_cast(mantissa), -9); + } + if (exponent == 0x0f && mantissa == 0x07) { + return CUDART_NAN_F; + } + return sign * ldexpf(1.0f + static_cast(mantissa) * 0.125f, + exponent - 7); +} + +__global__ void DequantizeLinearFp8Bf16Kernel(const uint8_t* input, + const uint8_t* zero_point, + const float* scales, uint16_t* output, + int64_t count, int64_t axis_size, + int64_t inner_size, + bool scalar_zero_point) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= count) return; + float value = DecodeE4M3FN(input[index]); + const int64_t scale_index = axis_size == 1 ? 0 : (index / inner_size) % axis_size; + if (zero_point != nullptr) { + value -= DecodeE4M3FN(zero_point[scalar_zero_point ? 0 : scale_index]); + } + value *= scales[scale_index]; + output[index] = __bfloat16_as_ushort(__float2bfloat16_rn(value)); +} + +__global__ void UntileHcBf16Kernel(const uint16_t* input, uint16_t* output, int kind) { + constexpr int64_t kVectorCount = 204800; + constexpr int64_t kElementCount = kVectorCount * 16; + const int64_t element = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (element >= kElementCount) return; + const int64_t vector_index = element / 16; + const int64_t lane = element % 16; + int64_t destination_vector; + if (kind == 0) { + const int64_t rp = vector_index & 127; + const int64_t t1 = vector_index >> 7; + const int64_t i = t1 % 10; + const int64_t t2 = t1 / 10; + const int64_t j = t2 % 20; + const int64_t s = t2 / 20; + destination_vector = (((s * 10 + i) * 128 + rp) * 20 + j); + } else { + const int64_t c = vector_index & 3; + const int64_t t1 = vector_index >> 2; + const int64_t rp = t1 & 31; + const int64_t t2 = t1 >> 5; + const int64_t i = t2 % 10; + const int64_t hb = t2 / 10; + destination_vector = (((i * 32 + rp) * 4 + c) * 160 + hb); + } + output[destination_vector * 16 + lane] = input[element]; +} + +} // namespace + +Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_type, + int64_t rows, int64_t width, int64_t* output, + float* partial_values, int64_t* partial_indices) { + ORT_RETURN_IF_NOT(rows >= 0 && width > 0, + "ArgMaxLastAxis requires a non-negative row count and positive vocabulary dimension"); + if (rows == 0) return Status::OK(); + switch (input_type) { + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, + partial_values, partial_indices); + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: + return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, + partial_values, partial_indices); + case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: + return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, + partial_values, partial_indices); + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "ArgMaxLastAxis input type is not supported"); + } +} + +Status LaunchEmbeddingRMSNorm(cudaStream_t stream, const int64_t* ids, const uint16_t* table, + const float* weight, float epsilon, uint16_t* output, + int64_t rows, int64_t vocab_size, int64_t hidden_size) { + ORT_RETURN_IF_NOT(hidden_size > 0 && vocab_size > 0, + "EmbeddingRMSNorm requires non-empty table dimensions"); + ORT_RETURN_IF_NOT(rows <= std::numeric_limits::max(), + "EmbeddingRMSNorm has too many rows for a CUDA grid"); + if (rows == 0) return Status::OK(); + EmbedNormKernel<<(rows), 512, 0, stream>>>( + ids, table, weight, epsilon, output, rows, hidden_size); + return CUDA_CALL(cudaGetLastError()); +} + +Status LaunchRowSelect(cudaStream_t stream, const int64_t* selection, + const uint16_t* const* inputs, uint16_t* const* outputs, + const int64_t* row_sizes) { + const dim3 grid(8, 3); + RowSelectKernel<<>>( + selection, inputs[0], inputs[1], inputs[2], outputs[0], outputs[1], outputs[2], + row_sizes[0], row_sizes[1], row_sizes[2]); + return CUDA_CALL(cudaGetLastError()); +} + +Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, + const uint16_t* projection, uint16_t* output, int64_t count, + int head_dim, int heads, int q_stride, + int64_t token_stride, int64_t q_offset, + int64_t gate_row_width, const uint16_t* residual, bool sigmoid_in_fp32, + bool silu_gate) { + if (count == 0) return Status::OK(); + const int64_t block_count = (count - 1) / kThreads + 1; + ORT_RETURN_IF_NOT(block_count <= std::numeric_limits::max(), + "SigmoidGateMul requires too many CUDA blocks"); +#define LAUNCH_GATE(BROADCAST, RESIDUAL, SILU) \ + SigmoidGateMulKernel<<(block_count), kThreads, 0, stream>>>( \ + reinterpret_cast(attention), \ + reinterpret_cast(projection), \ + reinterpret_cast<__nv_bfloat16*>(output), count, head_dim, heads, q_stride, \ + token_stride, q_offset, gate_row_width, reinterpret_cast(residual), sigmoid_in_fp32) + if (silu_gate) { + if (residual != nullptr) { + LAUNCH_GATE(false, true, true); + } else { + LAUNCH_GATE(false, false, true); + } + } else if (gate_row_width > 0) { + if (residual != nullptr) { + LAUNCH_GATE(true, true, false); + } else { + LAUNCH_GATE(true, false, false); + } + } else if (residual != nullptr) { + LAUNCH_GATE(false, true, false); + } else { + LAUNCH_GATE(false, false, false); + } +#undef LAUNCH_GATE + return CUDA_CALL(cudaGetLastError()); +} + +Status LaunchDequantizeLinearFp8Bf16(cudaStream_t stream, const uint8_t* input, + const uint8_t* zero_point, const float* scales, + uint16_t* output, int64_t count, + int64_t axis_size, int64_t inner_size, + bool scalar_zero_point) { + if (count == 0) return Status::OK(); + const int64_t block_count = (count - 1) / kThreads + 1; + ORT_RETURN_IF_NOT(block_count <= std::numeric_limits::max(), + "DequantizeLinear FP8/BF16 requires too many CUDA blocks"); + DequantizeLinearFp8Bf16Kernel<<(block_count), kThreads, 0, stream>>>( + input, zero_point, scales, output, count, axis_size, inner_size, scalar_zero_point); + return CUDA_CALL(cudaGetLastError()); +} + +Status LaunchUntileHyperConnectionWeights(cudaStream_t stream, const uint16_t* input, + uint16_t* output, int kind) { + constexpr int64_t kElementCount = 204800ll * 16; + constexpr int threads = 256; + constexpr int blocks = static_cast((kElementCount + threads - 1) / threads); + UntileHcBf16Kernel<<>>(input, output, kind); + return CUDA_CALL(cudaGetLastError()); +} + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h new file mode 100644 index 0000000000000..7cc6b05922159 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h @@ -0,0 +1,97 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime::contrib::cuda { + +Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_type, + int64_t rows, int64_t width, int64_t* output, + float* partial_values, int64_t* partial_indices); +inline int64_t GetArgMaxLastAxisPartitions(int64_t width) { + return 1 + (width - 1) / 4096; +} +Status LaunchEmbeddingRMSNorm(cudaStream_t stream, const int64_t* ids, const uint16_t* table, + const float* weight, float epsilon, uint16_t* output, + int64_t rows, int64_t vocab_size, int64_t hidden_size); +Status LaunchRowSelect(cudaStream_t stream, const int64_t* selection, + const uint16_t* const* inputs, uint16_t* const* outputs, + const int64_t* row_sizes); +Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, + const uint16_t* projection, uint16_t* output, int64_t count, + int head_dim, int heads, int q_stride, + int64_t token_stride, int64_t q_offset, + int64_t gate_row_width, const uint16_t* residual, bool sigmoid_in_fp32, + bool silu_gate); +Status LaunchDequantizeLinearFp8Bf16(cudaStream_t stream, const uint8_t* input, + const uint8_t* zero_point, const float* scales, + uint16_t* output, int64_t count, + int64_t axis_size, int64_t inner_size, + bool scalar_zero_point); +Status LaunchUntileHyperConnectionWeights(cudaStream_t stream, const uint16_t* input, + uint16_t* output, int kind); + +template +class ArgMaxLastAxis final : public onnxruntime::cuda::CudaKernel { + public: + explicit ArgMaxLastAxis(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + bool keepdims_; +}; + +class EmbeddingRMSNorm final : public onnxruntime::cuda::CudaKernel { + public: + explicit EmbeddingRMSNorm(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + float epsilon_; +}; + +class RowSelect final : public onnxruntime::cuda::CudaKernel { + public: + explicit RowSelect(const OpKernelInfo& info) : CudaKernel(info) {} + Status ComputeInternal(OpKernelContext* context) const override; +}; + +class SigmoidGateMul final : public onnxruntime::cuda::CudaKernel { + public: + explicit SigmoidGateMul(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int head_dim_; + int heads_; + int q_stride_; + int64_t tok_stride_; + int64_t q_offset_; + bool broadcast_gate_; + bool sigmoid_in_fp32_; + bool silu_gate_; +}; + +class DequantizeLinearFp8Bf16 final : public onnxruntime::cuda::CudaKernel { + public: + explicit DequantizeLinearFp8Bf16(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int64_t axis_; +}; + +class UntileHyperConnectionWeights final : public onnxruntime::cuda::CudaKernel { + public: + explicit UntileHyperConnectionWeights(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int kind_; +}; + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc new file mode 100644 index 0000000000000..982698c00b222 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc @@ -0,0 +1,165 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h" + +#include + +#include "contrib_ops/cpu/hyper_connection_helper.h" +#include "contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h" +#include "core/providers/cuda/cuda_type_conversion.h" + +namespace onnxruntime::contrib::cuda { + +#define REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + HyperConnectionLowRankMix, kMSDomain, 1, T, kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + HyperConnectionLowRankMix); + +REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(float) +REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(MLFloat16) +REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(BFloat16) + +#undef REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL + +template +HyperConnectionLowRankMix::HyperConnectionLowRankMix(const OpKernelInfo& info) + : CudaKernel(info), + num_branches_(info.GetAttrOrDefault("num_branches", 0)), + epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), + reduction_scale_(info.GetAttrOrDefault("reduction_scale", 1.0f)), + silu_alpha_(info.GetAttrOrDefault("silu_alpha", 1.0f)), + gate_div_(info.GetAttrOrDefault("gate_div", 1.0f)), + gate_mul_(info.GetAttrOrDefault("gate_mul", 1.0f)) { + ORT_ENFORCE(num_branches_ > 0, "HyperConnectionLowRankMix requires a positive num_branches"); + ORT_ENFORCE(gate_div_ != 0.0f, "gate_div must be non-zero"); +} + +template +Status HyperConnectionLowRankMix::ComputeInternal(OpKernelContext* context) const { + using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; + + const Tensor* streams = context->Input(0); + const Tensor* norm_scale = context->Input(1); + const Tensor* down_weight = context->Input(2); + const Tensor* up_weight = context->Input(3); + const Tensor* inject_weight = context->Input(4); + const Tensor* prev_block_output = context->Input(5); + const Tensor* prev_post_gate = context->Input(6); + + hyper_connection::StreamShape params; + ORT_RETURN_IF_ERROR(hyper_connection::ResolveStreamShape(streams->Shape(), num_branches_, params)); + ORT_RETURN_IF_NOT(params.branches <= std::numeric_limits::max() && + params.hidden <= std::numeric_limits::max() / params.branches, + "branch or hidden dimension is too large for CUDA"); + const int64_t branches = params.branches; + const int64_t hidden = params.hidden; + const int64_t features = branches * hidden; + + ORT_RETURN_IF_NOT(norm_scale->Shape().Size() == features && norm_scale->Shape().NumDimensions() == 1, + "norm_scale must have shape (num_branches * hidden)"); + ORT_RETURN_IF_NOT(down_weight->Shape().NumDimensions() == 2 && down_weight->Shape()[0] == features, + "down_weight must have shape (num_branches * hidden, rank)"); + const int64_t rank = down_weight->Shape()[1]; + ORT_RETURN_IF_NOT(rank > 0 && rank <= std::numeric_limits::max(), "rank must be positive"); + ORT_RETURN_IF_NOT(up_weight->Shape().NumDimensions() == 2 && up_weight->Shape()[0] == rank && + up_weight->Shape()[1] == features, + "up_weight must have shape (rank, num_branches * hidden)"); + ORT_RETURN_IF_NOT(inject_weight->Shape().NumDimensions() == 2 && inject_weight->Shape()[0] == features && + inject_weight->Shape()[1] == branches, + "inject_weight must have shape (num_branches * hidden, num_branches)"); + ORT_RETURN_IF_NOT((prev_block_output == nullptr) == (prev_post_gate == nullptr), + "prev_block_output and prev_post_gate must be provided together"); + + const bool has_prologue = prev_block_output != nullptr; + if (has_prologue) { + const auto& dims = streams->Shape().GetDims(); + TensorShapeVector expected_output(dims.begin(), dims.begin() + params.prefix_rank); + expected_output.push_back(hidden); + TensorShapeVector expected_gate(dims.begin(), dims.begin() + params.prefix_rank); + expected_gate.push_back(branches); + ORT_RETURN_IF_NOT(prev_block_output->Shape() == TensorShape(expected_output), + "prev_block_output must have shape (..., hidden)"); + ORT_RETURN_IF_NOT(prev_post_gate->Shape() == TensorShape(expected_gate), + "prev_post_gate must have shape (..., num_branches)"); + } + + Tensor* pre_output = context->Output(0, TensorShape(params.reduced_shape)); + TensorShapeVector gate_shape(params.reduced_shape.begin(), params.reduced_shape.end() - 1); + gate_shape.push_back(branches); + Tensor* post_gate = context->Output(1, TensorShape(gate_shape)); + Tensor* streams_out = context->Output(2, streams->Shape()); + + if (params.rows == 0) { + return Status::OK(); + } + + HyperConnectionLowRankMixConfig config; + config.rows = params.rows; + config.branches = static_cast(branches); + config.hidden = static_cast(hidden); + config.rank = static_cast(rank); + config.epsilon = epsilon_; + config.reduction_scale = reduction_scale_; + config.silu_alpha = silu_alpha_; + config.gate_div = gate_div_; + config.gate_mul = gate_mul_; + + IAllocatorUniquePtr streams_scratch; + HyperConnectionLowRankMixTensors tensors; + tensors.streams = reinterpret_cast(streams->Data()); + tensors.norm_scale = reinterpret_cast(norm_scale->Data()); + tensors.down_weight = reinterpret_cast(down_weight->Data()); + tensors.up_weight = reinterpret_cast(up_weight->Data()); + tensors.inject_weight = reinterpret_cast(inject_weight->Data()); + if (has_prologue) { + tensors.prev_block_output = reinterpret_cast(prev_block_output->Data()); + tensors.prev_post_gate = reinterpret_cast(prev_post_gate->Data()); + // The kernels need the post-mixed streams; use scratch memory if the graph does not consume them. + if (streams_out != nullptr) { + tensors.streams_out = reinterpret_cast(streams_out->MutableData()); + } else { + streams_scratch = GetScratchBuffer(static_cast(params.rows * features), GetComputeStream(context)); + tensors.streams_out = reinterpret_cast(streams_scratch.get()); + } + } else if (streams_out != nullptr) { + // Without a prologue the streams pass through unchanged. + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(streams_out->MutableDataRaw(), streams->DataRaw(), + streams->SizeInBytes(), cudaMemcpyDeviceToDevice, Stream(context))); + } + tensors.pre_output = reinterpret_cast(pre_output->MutableData()); + tensors.post_gate = reinterpret_cast(post_gate->MutableData()); + + // The fused kernels use 16-byte vector accesses on every tensor. + auto aligned16 = [](const void* ptr) { return ptr == nullptr || reinterpret_cast(ptr) % 16 == 0; }; + const bool aligned = aligned16(tensors.streams) && aligned16(tensors.norm_scale) && + aligned16(tensors.down_weight) && aligned16(tensors.up_weight) && + aligned16(tensors.inject_weight) && aligned16(tensors.prev_block_output) && + aligned16(tensors.streams_out); + + const auto& device_prop = GetDeviceProp(); + if (aligned && CanUseFusedHyperConnectionLowRankMix(device_prop, config)) { + auto scratch = GetScratchBuffer(HyperConnectionLowRankMixFusedScratchElements(config), + GetComputeStream(context)); + return LaunchHyperConnectionLowRankMixFused(device_prop, Stream(context), config, tensors, + scratch.get()); + } + + const size_t rows = static_cast(params.rows); + auto hn = GetScratchBuffer(rows * static_cast(features), GetComputeStream(context)); + auto down = GetScratchBuffer(rows * static_cast(rank), GetComputeStream(context)); + auto up = GetScratchBuffer(rows * static_cast(features), GetComputeStream(context)); + auto inject = GetScratchBuffer(rows * static_cast(branches), GetComputeStream(context)); + return LaunchHyperConnectionLowRankMixUnfused( + GetCublasHandle(context), Stream(context), config, tensors, reinterpret_cast(hn.get()), + reinterpret_cast(down.get()), reinterpret_cast(up.get()), + reinterpret_cast(inject.get())); +} + +template class HyperConnectionLowRankMix; +template class HyperConnectionLowRankMix; +template class HyperConnectionLowRankMix; + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h new file mode 100644 index 0000000000000..246f44afb4f8b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h @@ -0,0 +1,25 @@ +// 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 { + +template +class HyperConnectionLowRankMix final : public onnxruntime::cuda::CudaKernel { + public: + explicit HyperConnectionLowRankMix(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int64_t num_branches_; + float epsilon_; + float reduction_scale_; + float silu_alpha_; + float gate_div_; + float gate_mul_; +}; + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu new file mode 100644 index 0000000000000..25f9ffccc4194 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu @@ -0,0 +1,888 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Fused low-rank hyper-connection block (see the HyperConnectionLowRankMix schema in bert_defs.cc). +// +// Fused path (rows <= kHyperConnectionLowRankMixMaxFusedRows, 16-bit T, 16-byte aligned shapes): two kernels. +// HcDownKernel grid (column tiles + 1, splits, row groups) +// [post-mix of the previous block] + branchwise RMS norm + down GEMV + inject GEMV. +// The K = C * H reduction is split in slices that each stay inside one branch; every CTA +// writes fp32 partial sums for its slice. No atomics, no clusters, no grid barrier. +// HcUpPreKernel grid (ceil(H / 16), 1, row groups) +// reduces the partials in a fixed order, then ScaledSiLU, up GEMV, sigmoid, pre-mix and the +// inject gate. +// Both kernels are launched with programmatic dependent launch (sm_90+, plain launches elsewhere). Each one +// triggers its dependents immediately and prefetches its weights (which never depend on a previous kernel) +// into shared memory before griddepcontrol.wait; everything that consumes another kernel's output comes after +// the wait, so correctness never depends on the trigger position. +// Every intermediate is rounded to T at the points where the unfused ONNX subgraph rounds; the GEMVs +// accumulate in fp32 in a fixed (deterministic) order. +// +// Unfused path (any shape, any T): norm kernel + cuBLAS GEMMs (fp32 accumulation) + elementwise kernels. + +#include "contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h" + +#include +#include +#include +#include +#include +#include + +#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime::contrib::cuda { +namespace { + +constexpr int kThreads = 256; +constexpr int kColTile = 16; // down columns per CTA of the down kernel (two 16-byte weight vectors) +constexpr int kUpTile = 16; // features per CTA of the up kernel +constexpr int kMaxSliceRows = 1280; // K rows handled by one CTA of the down kernel +constexpr int kWeightSlots = 10; // 16-byte weight vectors per thread held in shared memory +constexpr int kMaxVecIters = 2; // 16-byte vectors per thread in the norm pass +constexpr int kMaxFusedHidden = kMaxVecIters * kThreads * 8; +constexpr int kMaxFusedRank = 512; +constexpr int kMaxFusedBranches = 8; + +__device__ __forceinline__ float SigmoidF32(float v) { + if (v >= 0.0f) return 1.0f / (1.0f + expf(-v)); + const float e = expf(v); + return e / (1.0f + e); +} + +template +__device__ __forceinline__ float Rnd(float v) { + return to_float(from_float(v)); +} + +template +__device__ __forceinline__ float BitsToFloat(uint16_t bits) { + if constexpr (std::is_same_v) { + return __uint_as_float(static_cast(bits) << 16); + } else { + return __half2float(__ushort_as_half(bits)); + } +} + +template +__device__ __forceinline__ uint16_t FloatToBits(float v) { + if constexpr (std::is_same_v) { + return __bfloat16_as_ushort(__float2bfloat16(v)); + } else { + return __half_as_ushort(__float2half(v)); + } +} + +template +__device__ __forceinline__ void Unpack2(uint32_t u, float& a, float& b) { + a = BitsToFloat(static_cast(u & 0xffffu)); + b = BitsToFloat(static_cast(u >> 16)); +} + +template +__device__ __forceinline__ uint32_t Pack2(float a, float b) { + return static_cast(FloatToBits(a)) | (static_cast(FloatToBits(b)) << 16); +} + +template +__device__ __forceinline__ void Unpack8(const uint4& u, float (&f)[8]) { + Unpack2(u.x, f[0], f[1]); + Unpack2(u.y, f[2], f[3]); + Unpack2(u.z, f[4], f[5]); + Unpack2(u.w, f[6], f[7]); +} + +template +__device__ __forceinline__ uint4 Pack8(const float (&f)[8]) { + uint4 u; + u.x = Pack2(f[0], f[1]); + u.y = Pack2(f[2], f[3]); + u.z = Pack2(f[4], f[5]); + u.w = Pack2(f[6], f[7]); + return u; +} + +__device__ __forceinline__ void GridWait() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void GridTrigger() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void CpAsync16(void* smem, const void* gmem) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(static_cast(__cvta_generic_to_shared(smem))), + "l"(gmem) + : "memory"); +#else + *reinterpret_cast(smem) = *reinterpret_cast(gmem); +#endif +} + +__device__ __forceinline__ void CpCommit() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.commit_group;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void CpWaitAll() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.wait_group 0;" ::: "memory"); +#endif +} + +struct DownParams { + int64_t rows; + int branches; + int hidden; + int rank; + int pieces; // slices per branch + int slice; // K rows per slice (multiple of 8) + int splits; // branches * pieces + int partial_stride; // rank + branches, rounded up to a multiple of 4 + float epsilon; +}; + +struct UpParams { + int64_t rows; + int branches; + int hidden; + int rank; + int splits; + int partial_stride; + float silu_alpha; + float reduction_scale; + float gate_div; + float gate_mul; +}; + +// CTA (j, s, z): K slice s (rows [k0, k0 + len) of the flattened stream, inside branch c) of the ROWS rows +// z * ROWS ... z * ROWS + ROWS - 1, which share every weight read. +// j < ceil(R / 16) owns down columns [16 j, 16 j + 16); the last j owns the C inject columns. +template +__global__ void __launch_bounds__(kThreads) + HcDownKernel(const T* __restrict__ streams, const T* __restrict__ prev_out, const T* __restrict__ prev_gate, + const T* __restrict__ norm_w, const T* __restrict__ w_down, const T* __restrict__ w_inject, + T* __restrict__ streams_out, float* __restrict__ partial, float* __restrict__ inv_rms_out, + DownParams p) { + const int j = blockIdx.x; + const int s = blockIdx.y; + const int64_t row0 = static_cast(blockIdx.z) * ROWS; + const int tid = threadIdx.x; + const int branches = p.branches; + const int hidden = p.hidden; + const int64_t K = static_cast(branches) * hidden; + const int c = s / p.pieces; + const int piece = s - c * p.pieces; + const int h0 = piece * p.slice; + const int len = min(p.slice, hidden - h0); + const int col_tiles = (p.rank + kColTile - 1) / kColTile; + const bool is_inject = (j == col_tiles); + const int64_t k0 = static_cast(c) * hidden + h0; + int64_t row_idx[ROWS]; + bool row_ok[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + row_ok[r] = row0 + r < p.rows; + row_idx[r] = row_ok[r] ? row0 + r : p.rows - 1; // keep loads of unused rows in bounds + } + + __shared__ __align__(16) uint4 wsm[kWeightSlots * kThreads]; + __shared__ __align__(16) uint16_t xs[ROWS][kMaxSliceRows]; + __shared__ float red[8][ROWS][kColTile]; + __shared__ float sq_red[8][ROWS]; + + const int cg = tid & 1; + const int rp = tid >> 1; + const int col = j * kColTile + cg * 8; + // Let the up kernel start right away and prefetch its weights; its results are only read after GridWait(). + GridTrigger(); + + // Weights do not depend on any previous kernel: start streaming them into shared memory first. + if (!is_inject) { + if (col < p.rank) { + const T* base = w_down + k0 * p.rank + col; +#pragma unroll + for (int i = 0; i < kWeightSlots; ++i) { + const int row = rp + 128 * i; + if (row < len) CpAsync16(&wsm[i * kThreads + tid], base + static_cast(row) * p.rank); + } + } + } else { + const int chunks = len * branches / 8; + for (int q = tid; q < chunks; q += kThreads) CpAsync16(&wsm[q], w_inject + k0 * branches + q * 8); + } + CpCommit(); + + // Norm weights for the 8 consecutive features of the slice this thread normalizes below. + const bool norm_active = tid * 8 < len; + uint4 wv = make_uint4(0u, 0u, 0u, 0u); + if (norm_active) wv = *reinterpret_cast(norm_w + k0 + tid * 8); + + // Everything below consumes results of the previous kernel. + GridWait(); + + const int nvec = hidden / 8; + uint4 sv[ROWS][kMaxVecIters]; + uint4 bv[ROWS][kMaxVecIters]; + float gate_c[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const uint4* sp = reinterpret_cast(streams + row_idx[r] * K + static_cast(c) * hidden); +#pragma unroll + for (int i = 0; i < kMaxVecIters; ++i) { + const int qi = tid + kThreads * i; + sv[r][i] = qi < nvec ? sp[qi] : make_uint4(0u, 0u, 0u, 0u); + } + if constexpr (kPost) { + gate_c[r] = to_float(prev_gate[row_idx[r] * branches + c]); + const uint4* bp = reinterpret_cast(prev_out + row_idx[r] * hidden); +#pragma unroll + for (int i = 0; i < kMaxVecIters; ++i) { + const int qi = tid + kThreads * i; + bv[r][i] = qi < nvec ? bp[qi] : make_uint4(0u, 0u, 0u, 0u); + } + } + } + + float sum_sq[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + sum_sq[r] = 0.0f; +#pragma unroll + for (int i = 0; i < kMaxVecIters; ++i) { + const int qi = tid + kThreads * i; + if (qi < nvec) { + const int h = qi * 8; + float v[8]; + Unpack8(sv[r][i], v); + if constexpr (kPost) { + float b[8]; + Unpack8(bv[r][i], b); +#pragma unroll + for (int m = 0; m < 8; ++m) v[m] = Rnd(fmaf(gate_c[r], b[m], v[m])); + } +#pragma unroll + for (int m = 0; m < 8; ++m) sum_sq[r] += v[m] * v[m]; + if (h >= h0 && h < h0 + len) { + const uint4 packed = Pack8(v); + *reinterpret_cast(&xs[r][h - h0]) = packed; + if constexpr (kPost) { + if (j == 0 && row_ok[r]) { + *reinterpret_cast(streams_out + row_idx[r] * K + static_cast(c) * hidden + h) = + packed; + } + } + } + } + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) sum_sq[r] += __shfl_xor_sync(0xffffffffu, sum_sq[r], off); + if ((tid & 31) == 0) sq_red[tid >> 5][r] = sum_sq[r]; + } + __syncthreads(); + float inv_rms[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + float total = 0.0f; +#pragma unroll + for (int w = 0; w < 8; ++w) total += sq_red[w][r]; + inv_rms[r] = rsqrtf(total / static_cast(hidden) + p.epsilon); + if (j == 0 && piece == 0 && tid == 0 && row_ok[r]) inv_rms_out[row_idx[r] * branches + c] = inv_rms[r]; + if (norm_active) { + float x[8], w[8]; + Unpack8(*reinterpret_cast(&xs[r][tid * 8]), x); + Unpack8(wv, w); +#pragma unroll + for (int m = 0; m < 8; ++m) x[m] = x[m] * inv_rms[r] * w[m]; + *reinterpret_cast(&xs[r][tid * 8]) = Pack8(x); + } + } + CpWaitAll(); + __syncthreads(); + + if (!is_inject) { + float acc[ROWS][8]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < 8; ++m) acc[r][m] = 0.0f; + } + if (col < p.rank) { +#pragma unroll + for (int i = 0; i < kWeightSlots; ++i) { + const int row = rp + 128 * i; + if (row < len) { + float wf[8]; + Unpack8(wsm[i * kThreads + tid], wf); +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float xv = BitsToFloat(xs[r][row]); +#pragma unroll + for (int m = 0; m < 8; ++m) acc[r][m] = fmaf(xv, wf[m], acc[r][m]); + } + } + } + } + // Lanes of a warp share cg = lane & 1 and differ in rp: reduce over lane bits 1..4. +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < 8; ++m) { +#pragma unroll + for (int off = 2; off < 32; off <<= 1) acc[r][m] += __shfl_xor_sync(0xffffffffu, acc[r][m], off); + } + if ((tid & 31) < 2) { +#pragma unroll + for (int m = 0; m < 8; ++m) red[tid >> 5][r][cg * 8 + m] = acc[r][m]; + } + } + __syncthreads(); + if (tid < kColTile * ROWS) { + const int r = tid / kColTile; + const int t = tid % kColTile; + float v = 0.0f; +#pragma unroll + for (int q = 0; q < 8; ++q) v += red[q][r][t]; + const int out_col = j * kColTile + t; + if (out_col < p.rank && row_ok[r]) partial[(row_idx[r] * p.splits + s) * p.partial_stride + out_col] = v; + } + } else { + const T* wi = reinterpret_cast(wsm); + float a[ROWS][kMaxFusedBranches]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < kMaxFusedBranches; ++m) a[r][m] = 0.0f; + } +#pragma unroll + for (int i = 0; i < 5; ++i) { + const int k = tid + kThreads * i; + if (k < len) { + float wf[kMaxFusedBranches]; +#pragma unroll + for (int m = 0; m < kMaxFusedBranches; ++m) wf[m] = m < branches ? to_float(wi[k * branches + m]) : 0.0f; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float xv = BitsToFloat(xs[r][k]); +#pragma unroll + for (int m = 0; m < kMaxFusedBranches; ++m) a[r][m] = fmaf(xv, wf[m], a[r][m]); + } + } + } +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < kMaxFusedBranches; ++m) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) a[r][m] += __shfl_xor_sync(0xffffffffu, a[r][m], off); + if ((tid & 31) == 0) red[tid >> 5][r][m] = a[r][m]; + } + } + __syncthreads(); + if (tid < branches * ROWS) { + const int r = tid / branches; + const int t = tid % branches; + float v = 0.0f; +#pragma unroll + for (int q = 0; q < 8; ++q) v += red[q][r][t]; + if (row_ok[r]) partial[(row_idx[r] * p.splits + s) * p.partial_stride + p.rank + t] = v; + } + } +} + +// CTA (hb, 0, z) owns features [16 hb, 16 hb + 16) of every branch for ROWS rows. +template +__global__ void __launch_bounds__(kThreads) + HcUpPreKernel(const T* __restrict__ streams, const T* __restrict__ norm_w, const T* __restrict__ w_up, + const float* __restrict__ partial, const float* __restrict__ inv_rms_in, + T* __restrict__ pre_out, T* __restrict__ gate_out, UpParams p) { + const int hb = blockIdx.x; + const int64_t row0 = static_cast(blockIdx.z) * ROWS; + const int tid = threadIdx.x; + const int branches = BRANCHES == 0 ? p.branches : BRANCHES; + const int rank = RANK == 0 ? p.rank : RANK; + constexpr int kBranchStorage = BRANCHES == 0 ? kMaxFusedBranches : BRANCHES; + const int hidden = p.hidden; + const int64_t K = static_cast(branches) * hidden; + const int h0 = hb * kUpTile; + const int groups = 2 * branches; // (branch, 8-feature half) pairs; divides 32 + const int phases = kThreads / groups; // row phases + const int g = tid % groups; + const int c = g >> 1; + const int half_id = g & 1; + const int rp = tid / groups; + const int col = h0 + half_id * 8; + const bool valid = col < hidden; + const int chunk_rows = phases * kWeightSlots; + int64_t row_idx[ROWS]; + bool row_ok[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + row_ok[r] = row0 + r < p.rows; + row_idx[r] = row_ok[r] ? row0 + r : p.rows - 1; + } + + __shared__ __align__(16) uint4 wsm[kWeightSlots * kThreads]; + __shared__ __align__(16) float sf[ROWS][RANK == 0 ? kMaxFusedRank : RANK]; + __shared__ float sg[ROWS][kBranchStorage * kUpTile]; + + // Prefetch the first chunk of up weights before waiting for the down kernel. + const T* wbase = w_up + static_cast(c) * hidden + col; + if (valid) { +#pragma unroll + for (int i = 0; i < kWeightSlots; ++i) { + const int row = rp + phases * i; + if (row < rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); + } + } + CpCommit(); + + // The next block's down kernel may start (and prefetch its weights) while this kernel waits. + GridTrigger(); + GridWait(); + + float xv[ROWS][kBranchStorage], inv[ROWS][kBranchStorage], nw[kBranchStorage]; + const int hl = tid; // feature owned by threads < 16 in the final pre-mix + const bool owner = tid < kUpTile && (h0 + hl) < hidden; + if (owner) { +#pragma unroll + for (int b = 0; b < kBranchStorage; ++b) { + if (b < branches) { + nw[b] = to_float(norm_w[static_cast(b) * hidden + h0 + hl]); +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + xv[r][b] = to_float(streams[row_idx[r] * K + static_cast(b) * hidden + h0 + hl]); + inv[r][b] = inv_rms_in[row_idx[r] * branches + b]; + } + } + } + } + + // Fixed-order reduction of the K-split partial sums, then ScaledSiLU; four columns per thread. + const int nq = rank / 4; + const int cq = tid * 4; + float4 va[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) va[r] = make_float4(0.f, 0.f, 0.f, 0.f); + if (tid < nq) { + for (int sp0 = 0; sp0 < p.splits; sp0 += 8) { + float4 pq[8][ROWS]; +#pragma unroll + for (int u = 0; u < 8; ++u) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + pq[u][r] = sp0 + u < p.splits + ? __ldca(reinterpret_cast( + partial + (row_idx[r] * p.splits + sp0 + u) * p.partial_stride + cq)) + : make_float4(0.f, 0.f, 0.f, 0.f); + } + } +#pragma unroll + for (int u = 0; u < 8; ++u) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + va[r].x += pq[u][r].x; + va[r].y += pq[u][r].y; + va[r].z += pq[u][r].z; + va[r].w += pq[u][r].w; + } + } + } +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float v[4] = {va[r].x, va[r].y, va[r].z, va[r].w}; + float out[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const float z = Rnd(Rnd(v[m]) * p.silu_alpha); + out[m] = Rnd(z * Rnd(SigmoidF32(z))); + } + *reinterpret_cast(&sf[r][cq]) = make_float4(out[0], out[1], out[2], out[3]); + } + } + if (hb == 0 && tid < branches * ROWS) { + const int r = tid / branches; + const int t = tid % branches; + float v = 0.0f; + for (int sp = 0; sp < p.splits; ++sp) v += partial[(row_idx[r] * p.splits + sp) * p.partial_stride + rank + t]; + const float inject = Rnd(v); + const float u = Rnd(SigmoidF32(Rnd(inject / p.gate_div))); + if (row_ok[r]) gate_out[row_idx[r] * branches + t] = from_float(u * p.gate_mul); + } + __syncthreads(); + + float acc[ROWS][8]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < 8; ++m) acc[r][m] = 0.0f; + } + for (int chunk0 = 0; chunk0 < rank; chunk0 += chunk_rows) { + if (chunk0 > 0 && valid) { +#pragma unroll + for (int i = 0; i < kWeightSlots; ++i) { + const int row = chunk0 + rp + phases * i; + if (row < rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); + } + CpCommit(); + } + CpWaitAll(); + if (valid) { +#pragma unroll + for (int i = 0; i < kWeightSlots; ++i) { + const int row = chunk0 + rp + phases * i; + if (row < rank) { + float wf[8]; + Unpack8(wsm[i * kThreads + tid], wf); +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float sv = sf[r][row]; +#pragma unroll + for (int m = 0; m < 8; ++m) acc[r][m] = fmaf(sv, wf[m], acc[r][m]); + } + } + } + } + } + // Lanes of a warp share g = lane % groups and differ in rp. + for (int off = groups; off < 32; off <<= 1) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < 8; ++m) acc[r][m] += __shfl_xor_sync(0xffffffffu, acc[r][m], off); + } + } + __syncthreads(); // every thread is done with its weights: reuse that shared memory for the cross-warp sums + float* red = reinterpret_cast(wsm); // [8 warps][ROWS][groups * 8] + if ((tid & 31) < groups) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < 8; ++m) red[((tid >> 5) * ROWS + r) * (groups * 8) + g * 8 + m] = acc[r][m]; + } + } + __syncthreads(); + if (tid < groups * 8) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + float v = 0.0f; +#pragma unroll + for (int q = 0; q < 8; ++q) v += red[(q * ROWS + r) * (groups * 8) + tid]; + const int gg = tid >> 3; + const int local = (gg & 1) * 8 + (tid & 7); + sg[r][(gg >> 1) * kUpTile + local] = Rnd(SigmoidF32(Rnd(v))); + } + } + __syncthreads(); + if (owner) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + float sum = 0.0f; +#pragma unroll + for (int b = 0; b < kBranchStorage; ++b) { + if (b < branches) { + const float n = Rnd(xv[r][b] * inv[r][b] * nw[b]); + sum = fmaf(n, sg[r][b * kUpTile + hl], sum); + } + } + if (row_ok[r]) pre_out[row_idx[r] * hidden + h0 + hl] = from_float(sum * p.reduction_scale); + } + } +} + +// ----------------------------------------------------------------------------- unfused path kernels + +template +__global__ void NormKernel(const T* __restrict__ streams, const T* __restrict__ prev_out, + const T* __restrict__ prev_gate, const T* __restrict__ norm_w, T* __restrict__ hn, + T* __restrict__ streams_out, int branches, int hidden, float epsilon) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage storage; + __shared__ float inv_rms; + const int64_t group = blockIdx.x; + const int64_t row = group / branches; + const int c = static_cast(group - row * branches); + const int64_t base = group * hidden; + float gate_c = 0.0f; + if constexpr (kPost) gate_c = to_float(prev_gate[group]); + + float sum_sq = 0.0f; + for (int h = threadIdx.x; h < hidden; h += kThreads) { + float v = to_float(streams[base + h]); + if constexpr (kPost) { + v = Rnd(fmaf(gate_c, to_float(prev_out[row * hidden + h]), v)); + streams_out[base + h] = from_float(v); + } + sum_sq += v * v; + } + const float total = BlockReduce(storage).Sum(sum_sq); + if (threadIdx.x == 0) inv_rms = rsqrtf(total / static_cast(hidden) + epsilon); + __syncthreads(); + for (int h = threadIdx.x; h < hidden; h += kThreads) { + float v = to_float(streams[base + h]); + if constexpr (kPost) v = Rnd(fmaf(gate_c, to_float(prev_out[row * hidden + h]), v)); + hn[base + h] = from_float(v * inv_rms * to_float(norm_w[static_cast(c) * hidden + h])); + } +} + +template +__global__ void ScaledSiluKernel(T* __restrict__ x, int64_t count, float alpha) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count) return; + const float z = Rnd(to_float(x[i]) * alpha); + const float sig = Rnd(SigmoidF32(z)); + x[i] = from_float(z * sig); +} + +template +__global__ void PreMixKernel(const T* __restrict__ hn, const T* __restrict__ up, T* __restrict__ pre, + int64_t count, int branches, int hidden, float reduction_scale) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count) return; + const int64_t row = i / hidden; + const int h = static_cast(i - row * hidden); + const int64_t K = static_cast(branches) * hidden; + float sum = 0.0f; + for (int b = 0; b < branches; ++b) { + const int64_t idx = row * K + static_cast(b) * hidden + h; + const float gate = Rnd(SigmoidF32(to_float(up[idx]))); + sum = fmaf(to_float(hn[idx]), gate, sum); + } + pre[i] = from_float(sum * reduction_scale); +} + +template +__global__ void PostGateKernel(const T* __restrict__ inject, T* __restrict__ gate, int64_t count, float gate_div, + float gate_mul) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count) return; + const float t = Rnd(to_float(inject[i]) / gate_div); + const float u = Rnd(SigmoidF32(t)); + gate[i] = from_float(u * gate_mul); +} + +template +constexpr cudaDataType_t CudaDataType() { + if constexpr (std::is_same_v) { + return CUDA_R_32F; + } else if constexpr (std::is_same_v) { + return CUDA_R_16F; + } else { + return CUDA_R_16BF; + } +} + +// Row-major C[m, n] = A[m, k] * B[k, n] with fp32 accumulation and T output. +template +Status Gemm(cublasHandle_t cublas, int64_t m, int64_t n, int64_t k, const T* a, const T* b, T* c) { + ORT_RETURN_IF_NOT(m <= std::numeric_limits::max() && n <= std::numeric_limits::max() && + k <= std::numeric_limits::max(), + "HyperConnectionLowRankMix GEMM dimension is too large"); + const float alpha = 1.0f; + const float beta = 0.0f; + constexpr cudaDataType_t type = CudaDataType(); + CUBLAS_RETURN_IF_ERROR(cublasGemmEx(cublas, CUBLAS_OP_N, CUBLAS_OP_N, static_cast(n), static_cast(m), + static_cast(k), &alpha, b, type, static_cast(n), a, type, + static_cast(k), &beta, c, type, static_cast(n), + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + return Status::OK(); +} + +int64_t BlocksFor(int64_t count) { return (count + kThreads - 1) / kThreads; } + +struct FusedGeometry { + int pieces; + int slice; + int splits; + int col_tiles; + int partial_stride; +}; + +FusedGeometry GetFusedGeometry(const HyperConnectionLowRankMixConfig& config) { + FusedGeometry g; + g.pieces = (config.hidden + kMaxSliceRows - 1) / kMaxSliceRows; + g.slice = ((config.hidden + g.pieces - 1) / g.pieces + 7) / 8 * 8; + g.pieces = (config.hidden + g.slice - 1) / g.slice; + g.splits = config.branches * g.pieces; + g.col_tiles = (config.rank + kColTile - 1) / kColTile; + g.partial_stride = (config.rank + config.branches + 3) / 4 * 4; // keeps every row of partials 16-byte aligned + return g; +} + +template +Status LaunchWithProgrammaticSerialization(bool use_pdl, Kernel kernel, dim3 grid, dim3 block, cudaStream_t stream, + Args... args) { +#if defined(CUDART_VERSION) && CUDART_VERSION >= 12030 + if (use_pdl) { + cudaLaunchAttribute attribute; + attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute.val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t launch_config = {}; + launch_config.gridDim = grid; + launch_config.blockDim = block; + launch_config.dynamicSmemBytes = 0; + launch_config.stream = stream; + launch_config.attrs = &attribute; + launch_config.numAttrs = 1; + return CUDA_CALL(cudaLaunchKernelEx(&launch_config, kernel, args...)); + } +#else + ORT_UNUSED_PARAMETER(use_pdl); +#endif + kernel<<>>(args...); + return CUDA_CALL(cudaGetLastError()); +} + +template +Status LaunchFusedRows(bool use_pdl, cudaStream_t stream, const HyperConnectionLowRankMixConfig& config, + const HyperConnectionLowRankMixTensors& tensors, const FusedGeometry& g, float* partial, + float* inv_rms, const DownParams& down, const UpParams& up) { + const unsigned row_groups = static_cast((config.rows + ROWS - 1) / ROWS); + const dim3 down_grid(g.col_tiles + 1, g.splits, row_groups); + const dim3 up_grid((config.hidden + kUpTile - 1) / kUpTile, 1, row_groups); + if (tensors.prev_block_output != nullptr) { + ORT_RETURN_IF_ERROR(LaunchWithProgrammaticSerialization( + use_pdl, HcDownKernel, down_grid, dim3(kThreads), stream, tensors.streams, + tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, tensors.down_weight, + tensors.inject_weight, tensors.streams_out, partial, inv_rms, down)); + } else { + ORT_RETURN_IF_ERROR(LaunchWithProgrammaticSerialization( + use_pdl, HcDownKernel, down_grid, dim3(kThreads), stream, tensors.streams, + static_cast(nullptr), static_cast(nullptr), tensors.norm_scale, + tensors.down_weight, tensors.inject_weight, static_cast(nullptr), partial, inv_rms, down)); + } + + const T* up_streams = tensors.prev_block_output != nullptr ? tensors.streams_out : tensors.streams; + if (config.branches == 4 && config.rank == 320) { + return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), stream, + up_streams, tensors.norm_scale, tensors.up_weight, + static_cast(partial), + static_cast(inv_rms), tensors.pre_output, + tensors.post_gate, up); + } + return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), stream, + up_streams, tensors.norm_scale, tensors.up_weight, + static_cast(partial), + static_cast(inv_rms), tensors.pre_output, + tensors.post_gate, up); +} + +} // namespace + +template +bool CanUseFusedHyperConnectionLowRankMix(const cudaDeviceProp& device_prop, + const HyperConnectionLowRankMixConfig& config) { + if constexpr (sizeof(T) != 2) { + return false; + } else { + const bool branches_ok = config.branches == 1 || config.branches == 2 || config.branches == 4 || + config.branches == 8; + return device_prop.major >= 8 && config.rows >= 1 && + config.rows <= kHyperConnectionLowRankMixMaxFusedRows && branches_ok && + config.hidden % 8 == 0 && config.hidden <= kMaxFusedHidden && + config.rank % 8 == 0 && config.rank >= 8 && config.rank <= kMaxFusedRank; + } +} + +size_t HyperConnectionLowRankMixFusedScratchElements(const HyperConnectionLowRankMixConfig& config) { + const FusedGeometry g = GetFusedGeometry(config); + return static_cast(config.rows) * + (static_cast(g.splits) * g.partial_stride + static_cast(config.branches)); +} + +template +Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, cudaStream_t stream, + const HyperConnectionLowRankMixConfig& config, + const HyperConnectionLowRankMixTensors& tensors, + float* scratch) { + if constexpr (sizeof(T) != 2) { + ORT_UNUSED_PARAMETER(device_prop); + ORT_UNUSED_PARAMETER(stream); + ORT_UNUSED_PARAMETER(config); + ORT_UNUSED_PARAMETER(tensors); + ORT_UNUSED_PARAMETER(scratch); + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "fused HyperConnectionLowRankMix requires a 16-bit type"); + } else { + ORT_RETURN_IF_NOT(CanUseFusedHyperConnectionLowRankMix(device_prop, config), + "unsupported problem for the fused HyperConnectionLowRankMix path"); + const FusedGeometry g = GetFusedGeometry(config); + float* partial = scratch; + float* inv_rms = scratch + static_cast(config.rows) * g.splits * g.partial_stride; + + DownParams down{config.rows, config.branches, config.hidden, config.rank, g.pieces, g.slice, + g.splits, g.partial_stride, config.epsilon}; + UpParams up{config.rows, config.branches, config.hidden, config.rank, g.splits, g.partial_stride, + config.silu_alpha, config.reduction_scale, config.gate_div, config.gate_mul}; + const bool use_pdl = device_prop.major >= 9; + if (config.rows == 1) { + return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); + } + return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); + } +} + +template +Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_t stream, + const HyperConnectionLowRankMixConfig& config, + const HyperConnectionLowRankMixTensors& tensors, + T* hn, T* down, T* up, T* inject) { + const int64_t rows = config.rows; + const int64_t K = static_cast(config.branches) * config.hidden; + const int64_t groups = rows * config.branches; + ORT_RETURN_IF_NOT(groups <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); + + if (tensors.prev_block_output != nullptr) { + NormKernel<<(groups), kThreads, 0, stream>>>( + tensors.streams, tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, hn, + tensors.streams_out, config.branches, config.hidden, config.epsilon); + } else { + NormKernel<<(groups), kThreads, 0, stream>>>( + tensors.streams, nullptr, nullptr, tensors.norm_scale, hn, nullptr, config.branches, config.hidden, + config.epsilon); + } + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + + ORT_RETURN_IF_ERROR(Gemm(cublas, rows, config.rank, K, hn, tensors.down_weight, down)); + ORT_RETURN_IF_ERROR(Gemm(cublas, rows, config.branches, K, hn, tensors.inject_weight, inject)); + + const int64_t down_count = rows * config.rank; + ScaledSiluKernel<<(BlocksFor(down_count)), kThreads, 0, stream>>>( + down, down_count, config.silu_alpha); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + ORT_RETURN_IF_ERROR(Gemm(cublas, rows, K, config.rank, down, tensors.up_weight, up)); + + const int64_t pre_count = rows * config.hidden; + ORT_RETURN_IF_NOT(BlocksFor(pre_count) <= std::numeric_limits::max(), + "CUDA launch requires too many blocks"); + PreMixKernel<<(BlocksFor(pre_count)), kThreads, 0, stream>>>( + hn, up, tensors.pre_output, pre_count, config.branches, config.hidden, config.reduction_scale); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + PostGateKernel<<(BlocksFor(groups)), kThreads, 0, stream>>>( + inject, tensors.post_gate, groups, config.gate_div, config.gate_mul); + return CUDA_CALL(cudaGetLastError()); +} + +#define INSTANTIATE(T) \ + template bool CanUseFusedHyperConnectionLowRankMix(const cudaDeviceProp&, \ + const HyperConnectionLowRankMixConfig&); \ + template Status LaunchHyperConnectionLowRankMixFused( \ + const cudaDeviceProp&, cudaStream_t, const HyperConnectionLowRankMixConfig&, \ + const HyperConnectionLowRankMixTensors&, float*); \ + template Status LaunchHyperConnectionLowRankMixUnfused( \ + cublasHandle_t, cudaStream_t, const HyperConnectionLowRankMixConfig&, \ + const HyperConnectionLowRankMixTensors&, T*, T*, T*, T*); + +INSTANTIATE(float) +INSTANTIATE(half) +INSTANTIATE(__nv_bfloat16) + +#undef INSTANTIATE + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h new file mode 100644 index 0000000000000..c3c9c24c4842b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h @@ -0,0 +1,67 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include +#include + +#include "core/common/status.h" + +namespace onnxruntime::contrib::cuda { + +struct HyperConnectionLowRankMixConfig { + int64_t rows = 0; + int branches = 0; // C + int hidden = 0; // H, width of one branch; the streams have C * H features + int rank = 0; // R, width of the low-rank bottleneck + float epsilon = 1e-5f; + float reduction_scale = 1.0f; + float silu_alpha = 1.0f; + float gate_div = 1.0f; + float gate_mul = 1.0f; +}; + +template +struct HyperConnectionLowRankMixTensors { + const T* streams = nullptr; + const T* norm_scale = nullptr; + const T* down_weight = nullptr; + const T* up_weight = nullptr; + const T* inject_weight = nullptr; + const T* prev_block_output = nullptr; // prologue inputs; both null or both non-null + const T* prev_post_gate = nullptr; + T* pre_output = nullptr; + T* post_gate = nullptr; + T* streams_out = nullptr; // written when the prologue is present; must then be non-null +}; + +// Rows up to this value use the two-kernel fused path; larger inputs use the cuBLAS based path. +constexpr int64_t kHyperConnectionLowRankMixMaxFusedRows = 8; + +// True when the fused two-kernel path supports this problem on the given device. +template +bool CanUseFusedHyperConnectionLowRankMix(const cudaDeviceProp& device_prop, + const HyperConnectionLowRankMixConfig& config); + +// Number of float elements of scratch memory the fused path needs. +size_t HyperConnectionLowRankMixFusedScratchElements(const HyperConnectionLowRankMixConfig& config); + +template +Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, cudaStream_t stream, + const HyperConnectionLowRankMixConfig& config, + const HyperConnectionLowRankMixTensors& tensors, + float* scratch); + +// General path built from cuBLAS GEMMs and elementwise kernels. Scratch buffers (in elements of T): +// hn [rows, C * H], down [rows, R], up [rows, C * H], inject [rows, C]. +template +Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_t stream, + const HyperConnectionLowRankMixConfig& config, + const HyperConnectionLowRankMixTensors& tensors, + T* hn, T* down, T* up, T* inject); + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 546b5b12e178f..de29a3579e58d 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -170,6 +170,13 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, LinearAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, GatedDeltaNet); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, GatedDeltaNet); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, GatedDeltaNet); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, ArgMaxLastAxis); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, ArgMaxLastAxis); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, ArgMaxLastAxis); +class CUDA_MS_OP_CLASS_NAME(1, EmbeddingRMSNorm); +class CUDA_MS_OP_CLASS_NAME(1, RowSelect); +class CUDA_MS_OP_CLASS_NAME(1, SigmoidGateMul); +class CUDA_MS_OP_CLASS_NAME(1, UntileHyperConnectionWeights); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, LinearAttentionGate); @@ -188,6 +195,9 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, HyperConnectionPreMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, HyperConnectionPostMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, HyperConnectionPostMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, HyperConnectionPostMix); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, HyperConnectionLowRankMix); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, HyperConnectionLowRankMix); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, HyperConnectionLowRankMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, EngramGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, EngramGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, EngramGate); @@ -534,6 +544,13 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -552,6 +569,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index fecac2f17f523..5e9751c7d445a 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4578,11 +4578,12 @@ ONNX_MS_OPERATOR_SET_SCHEMA( 1.0e-5f) .Input(0, "key", - "Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size).", + "Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size) " + "or (total_tokens, hc_mult, hidden_size).", "T") .Input(1, "query", - "Hidden-state queries with shape (batch_size, sequence_length, hc_mult, hidden_size).", + "Hidden-state queries with the same shape as key.", "T") .Input(2, "value", @@ -4605,7 +4606,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Output(0, "output", - "Gated value tensor with shape (batch_size, sequence_length, hc_mult, hidden_size).", + "Gated value tensor with the same shape as key.", "T") .Output(1, "gated_value_normed", @@ -5624,6 +5625,147 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } })); +constexpr const char* FusedTensorOps_ver1_doc = R"DOC( +CUDA contrib operators for fused tensor operations used by transformer models. +)DOC"; + +ONNX_MS_OPERATOR_SET_SCHEMA( + ArgMaxLastAxis, 1, + OpSchema() + .SetDoc(FusedTensorOps_ver1_doc) + .Attr("keepdims", "Keep a singleton final axis when nonzero.", AttributeProto::INT, + static_cast(0)) + .Input(0, "logits", "Logits with vocabulary on the last axis.", "T") + .Output(0, "indices", "First maximum index for each leading row.", "I") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Supported logit types.") + .TypeConstraint("I", {"tensor(int64)"}, "Argmax indices are int64.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + updateOutputElemType(ctx, 0, ONNX_NAMESPACE::TensorProto_DataType_INT64); + if (!hasInputShape(ctx, 0)) return; + const auto& input_shape = getInputShape(ctx, 0); + if (input_shape.dim_size() < 1) { + fail_shape_inference("ArgMaxLastAxis: logits must have rank at least one"); + } + ONNX_NAMESPACE::TensorShapeProto output_shape; + for (int i = 0; i + 1 < input_shape.dim_size(); ++i) { + *output_shape.add_dim() = input_shape.dim(i); + } + if (getAttribute(ctx, "keepdims", static_cast(0)) != 0) { + output_shape.add_dim()->set_dim_value(1); + } + updateOutputShape(ctx, 0, output_shape); + })); + +ONNX_MS_OPERATOR_SET_SCHEMA( + EmbeddingRMSNorm, 1, + OpSchema() + .SetDoc(FusedTensorOps_ver1_doc) + .Attr("epsilon", "RMS normalization epsilon.", AttributeProto::FLOAT) + .Input(0, "ids", "Embedding row indices with any shape.", "I") + .Input(1, "table", "BF16 embedding matrix [vocabulary, hidden].", "T") + .Input(2, "weight", "FP32 RMS normalization weight [hidden].", "F") + .Output(0, "normalized", "BF16 normalized embeddings [shape(ids), hidden].", "T") + .TypeConstraint("I", {"tensor(int64)"}, "Indices use int64.") + .TypeConstraint("T", {"tensor(bfloat16)"}, "Embedding and result use BF16.") + .TypeConstraint("F", {"tensor(float)"}, "Normalization weight uses FP32.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 1, 0); + if (!hasInputShape(ctx, 0) || !hasInputShape(ctx, 1)) return; + const auto& ids_shape = getInputShape(ctx, 0); + const auto& table_shape = getInputShape(ctx, 1); + if (table_shape.dim_size() != 2) { + fail_shape_inference("EmbeddingRMSNorm: table must have rank two"); + } + ONNX_NAMESPACE::TensorShapeProto output_shape = ids_shape; + *output_shape.add_dim() = table_shape.dim(1); + updateOutputShape(ctx, 0, output_shape); + })); + +ONNX_MS_OPERATOR_SET_SCHEMA( + RowSelect, 1, + OpSchema() + .SetDoc(FusedTensorOps_ver1_doc) + .Input(0, "selection", "Scalar-like int64 row selector; zero selects row 0.", "I") + .Input(1, "x0", "BF16 tensor with shape [1, 2, ...].", "T") + .Input(2, "x1", "BF16 tensor with shape [1, 2, ...].", "T") + .Input(3, "x2", "BF16 tensor with shape [1, 2, ...].", "T") + .Output(0, "y0", "Selected BF16 tensor with shape [1, 1, ...].", "T") + .Output(1, "y1", "Selected BF16 tensor with shape [1, 1, ...].", "T") + .Output(2, "y2", "Selected BF16 tensor with shape [1, 1, ...].", "T") + .TypeConstraint("I", {"tensor(int64)"}, "Selector uses int64.") + .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and outputs use BF16.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + for (int i = 0; i < 3; ++i) { + propagateElemTypeFromInputToOutput(ctx, i + 1, i); + if (!hasInputShape(ctx, i + 1)) continue; + ONNX_NAMESPACE::TensorShapeProto output_shape = getInputShape(ctx, i + 1); + if (output_shape.dim_size() < 2) { + fail_shape_inference("RowSelect: inputs must have rank at least two"); + } + output_shape.mutable_dim(1)->set_dim_value(1); + updateOutputShape(ctx, i, output_shape); + } + })); + +ONNX_MS_OPERATOR_SET_SCHEMA( + SigmoidGateMul, 1, + OpSchema() + .SetDoc(FusedTensorOps_ver1_doc) + .Attr("head_dim", "Dimension of each attention head.", AttributeProto::INT, static_cast(0)) + .Attr("num_heads", "Number of attention heads.", AttributeProto::INT, static_cast(0)) + .Attr("q_stride", "Per-head stride in the packed projection.", AttributeProto::INT, static_cast(0)) + .Attr("tok_stride", "Per-token stride in the packed projection.", AttributeProto::INT, static_cast(0)) + .Attr("q_offset", "Offset of the query projection within each token.", AttributeProto::INT, static_cast(0)) + .Attr("broadcast_gate", + "When one, projection has the attention shape with final dimension one. " + "Packed layout attributes must be zero. CUDA only.", + AttributeProto::INT, static_cast(0)) + .Attr("sigmoid_in_fp32", + "One computes sigmoid in FP32 before BF16 conversion. Zero preserves " + "the input-type intermediate arithmetic of CUDA Sigmoid, including subtraction on negative inputs.", + AttributeProto::INT, static_cast(1)) + .Attr("silu_gate", + "When one, projection and attention have identical shapes. Multiply the " + "BF16-rounded sigmoid by projection and round to BF16 before multiplying attention. " + "Cannot be combined with broadcast_gate or packed layout attributes.", + AttributeProto::INT, static_cast(0)) + .Input(0, "attention", "BF16 attention values.", "T") + .Input(1, "projection", "BF16 packed projection containing gate values.", "T") + .Input(2, "residual", "Optional BF16 residual with the attention shape.", "T", OpSchema::Optional) + .Output(0, "output", + "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 " + "before optionally adding residual and rounding again.", + "T") + .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and output use BF16.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + if (hasInputShape(ctx, 0)) updateOutputShape(ctx, 0, getInputShape(ctx, 0)); + })); + +ONNX_MS_OPERATOR_SET_SCHEMA( + UntileHyperConnectionWeights, 1, + OpSchema() + .SetDoc(FusedTensorOps_ver1_doc) + .Attr("kind", "Untile layout kind, either 0 or 1.", AttributeProto::INT) + .Input(0, "tiled", "Tiled BF16 weight [160, 20480].", "T") + .Input(1, "unused", "Reserved optional input retained for exporter compatibility.", "T", + OpSchema::Optional) + .Output(0, "untiled", "Untiled BF16 weight.", "T") + .TypeConstraint("T", {"tensor(bfloat16)"}, "Input and output use BF16.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + if (!hasInputShape(ctx, 0)) return; + const int64_t kind = getAttribute(ctx, "kind", static_cast(-1)); + if (kind != 0 && kind != 1) { + fail_shape_inference("UntileHyperConnectionWeights: kind must be zero or one"); + } + ONNX_NAMESPACE::TensorShapeProto output_shape; + output_shape.add_dim()->set_dim_value(kind == 0 ? 10240 : 320); + output_shape.add_dim()->set_dim_value(kind == 0 ? 320 : 10240); + updateOutputShape(ctx, 0, output_shape); + })); + constexpr const char* LinearAttentionGate_ver1_doc = R"DOC( Fuses the gate projections that feed LinearAttention's gated-delta recurrence: @@ -5900,6 +6042,113 @@ ONNX_MS_OPERATOR_SET_SCHEMA( propagateShapeFromInputToOutput(ctx, 0, 0); })); +constexpr const char* HyperConnectionLowRankMix_ver1_doc = R"DOC( +One hyper-connection block with a low-rank gate network, fused into one operator. C = num_branches streams of +width H are kept as one flattened (..., C * H) tensor. With the optional prologue the operator first applies the +post-mix of the previous block (identity stream mix): + + S[..., c, h] = streams[..., c, h] + prev_post_gate[..., c] * prev_block_output[..., h] + +and otherwise S = streams. It then computes, with `rms` taken per branch over H: + + hn[..., c, h] = S[..., c, h] / sqrt(mean_h(S[..., c, :]^2) + epsilon) * norm_scale[c, h] + d = ScaledSiLU(hn @ down_weight, silu_alpha) (..., R) + g[..., c, h] = sigmoid(d @ up_weight) per stream and feature + pre_output = reduction_scale * sum_c(hn[..., c, h] * g[..., c, h]) (..., H) + post_gate = sigmoid((hn @ inject_weight) / gate_div) * gate_mul (..., C) + +pre_output is the input of the next transformer block and post_gate is prev_post_gate of the next call. +streams_out is S (the streams after the prologue; a copy of streams when there is no prologue). + +Numerics contract. Matrix products and reductions accumulate in float32. The result of every step is rounded +to T once, as the output of the equivalent standard operator would be: +S = cast_T(streams + prev_post_gate * prev_block_output), hn = cast_T(S * inv_rms * norm_scale), +down = cast_T(hn @ down_weight), Z = cast_T(silu_alpha * down), SG = cast_T(sigmoid(Z)), hidden = cast_T(Z * SG), +up = cast_T(hidden @ up_weight), g = cast_T(sigmoid(up)), pre_output = cast_T(reduction_scale * sum_c(hn * g)), +inject = cast_T(hn @ inject_weight), post_gate = cast_T(cast_T(sigmoid(cast_T(inject / gate_div))) * gate_mul). +sigmoid is evaluated in float32 with expf and rounded once to T; it does not reproduce a Sigmoid kernel that +rounds to T after every elementary operation. The accumulation order of the matrix products is fixed, so the +results are deterministic, but they are not bitwise identical to cuBLAS. +)DOC"; + +ONNX_MS_OPERATOR_SET_SCHEMA( + HyperConnectionLowRankMix, 1, + OpSchema() + .SetDoc(HyperConnectionLowRankMix_ver1_doc) + .Attr("num_branches", "Number of residual streams C. The last dimension of streams is C * H.", + AttributeProto::INT, true) + .Attr("epsilon", "Epsilon added to the mean square of each branch before the reciprocal square root.", + AttributeProto::FLOAT, 1e-5f) + .Attr("reduction_scale", "Multiplier applied to the sum over branches of pre_output.", + AttributeProto::FLOAT, 1.0f) + .Attr("silu_alpha", "Scale of the ScaledSiLU between the down and up projections.", + AttributeProto::FLOAT, 1.0f) + .Attr("gate_div", "Divisor applied to the inject logits before the sigmoid. Must not be zero.", + AttributeProto::FLOAT, 1.0f) + .Attr("gate_mul", "Multiplier applied to the sigmoid of the inject logits.", + AttributeProto::FLOAT, 1.0f) + .Input(0, "streams", + "Flattened streams (..., C * H). With the prologue these are the streams before the previous " + "block's post-mix.", + "T") + .Input(1, "norm_scale", "Per-branch RMS norm weight with shape (C * H).", "T") + .Input(2, "down_weight", "Down projection with shape (C * H, R), row-major [K, N].", "T") + .Input(3, "up_weight", "Up projection with shape (R, C * H), row-major [K, N].", "T") + .Input(4, "inject_weight", "Post-gate projection with shape (C * H, C), row-major [K, N].", "T") + .Input(5, "prev_block_output", + "Optional output of the previous block with shape (..., H). Provide together with prev_post_gate.", + "T", OpSchema::Optional) + .Input(6, "prev_post_gate", + "Optional post-mix gate of the previous block with shape (..., C). Provide together with " + "prev_block_output.", + "T", OpSchema::Optional) + .Output(0, "pre_output", "Block input with shape (..., H).", "T") + .Output(1, "post_gate", "Post-mix gate for the next call with shape (..., C).", "T") + .Output(2, "streams_out", + "Optional streams after the prologue, with the same shape as streams. Equal to streams when the " + "prologue inputs are absent.", + "T", OpSchema::Optional) + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain all floating-point inputs and outputs to the same type.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateElemTypeFromInputToOutput(ctx, 0, 1); + if (ctx.getNumOutputs() > 2) { + propagateElemTypeFromInputToOutput(ctx, 0, 2); + } + if (!hasInputShape(ctx, 0)) { + return; + } + const auto branches = getAttribute(ctx, "num_branches", 0); + if (branches < 1) { + fail_shape_inference("num_branches must be positive."); + } + const auto& input_shape = getInputShape(ctx, 0); + const int rank = input_shape.dim_size(); + if (rank < 1) { + fail_shape_inference("streams must have rank at least 1."); + } + auto* pre_shape = getOutputShape(ctx, 0); + auto* gate_shape = getOutputShape(ctx, 1); + for (int i = 0; i < rank - 1; ++i) { + *pre_shape->add_dim() = input_shape.dim(i); + *gate_shape->add_dim() = input_shape.dim(i); + } + auto* hidden = pre_shape->add_dim(); + const auto& last = input_shape.dim(rank - 1); + if (last.has_dim_value()) { + if (last.dim_value() % branches != 0) { + fail_shape_inference("The last dimension of streams must be divisible by num_branches."); + } + hidden->set_dim_value(last.dim_value() / branches); + } + gate_shape->add_dim()->set_dim_value(branches); + if (ctx.getNumOutputs() > 2) { + propagateShapeFromInputToOutput(ctx, 0, 2); + } + })); + constexpr const char* GatedAdd_ver1_doc = R"DOC( Adds one tensor to another tensor scaled by a per-row gate: diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index d68cdd392666b..56939b20c829b 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -94,8 +94,14 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, 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); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionLowRankMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNet); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNetStateReplay); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ArgMaxLastAxis); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, EmbeddingRMSNorm); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RowSelect); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SigmoidGateMul); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, UntileHyperConnectionWeights); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedRMSNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedAdd); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, NGramHashMapping); @@ -228,8 +234,14 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/test/contrib_ops/engram_ops_test.cc b/onnxruntime/test/contrib_ops/engram_ops_test.cc index 0471be6941f53..6f64d835d7e9e 100644 --- a/onnxruntime/test/contrib_ops/engram_ops_test.cc +++ b/onnxruntime/test/contrib_ops/engram_ops_test.cc @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -69,13 +70,15 @@ bool IsTypeSupported() { // Runs the tester on CUDA only for BFloat16, otherwise on the default set of execution providers. template -void RunOnSupportedProviders(OpTester& test) { +void RunOnSupportedProviders(OpTester& test, bool packed = false) { if constexpr (std::is_same_v) { std::vector> providers; providers.push_back(DefaultCudaExecutionProvider()); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } else { - test.Run(); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + packed ? std::unordered_set{kWebGpuExecutionProvider} + : std::unordered_set{}); } } @@ -151,7 +154,7 @@ void RunEngramGateTest(float tolerance) { // Verifies gated_value_normed: RMSNorm is applied independently to each hyper-connection branch // (each hidden_size-sized slice), not over the concatenated hc_mult * hidden_size dimension. template -void RunEngramGateNormedTest(float tolerance) { +void RunEngramGateNormedTest(float tolerance, bool packed = false) { if (!IsTypeSupported()) { GTEST_SKIP() << "No execution provider available for this type"; } @@ -204,17 +207,21 @@ void RunEngramGateNormedTest(float tolerance) { OpTester test("EngramGate", 1, kMSDomain); test.AddAttribute("epsilon", kEpsilon); - test.AddInput("key", {1, 1, hc_mult, hidden_size}, ToTensorType(key)); - test.AddInput("query", {1, 1, hc_mult, hidden_size}, ToTensorType(query)); - test.AddInput("value", {1, 1, hidden_size}, ToTensorType(value)); + const std::vector key_shape = packed ? std::vector{1, hc_mult, hidden_size} + : std::vector{1, 1, hc_mult, hidden_size}; + const std::vector value_shape = packed ? std::vector{1, hidden_size} + : std::vector{1, 1, hidden_size}; + test.AddInput("key", key_shape, ToTensorType(key)); + test.AddInput("query", key_shape, ToTensorType(query)); + test.AddInput("value", value_shape, ToTensorType(value)); test.AddInput("key_norm_scale", {hc_mult, hidden_size}, ToTensorType(key_scale)); test.AddInput("query_norm_scale", {hc_mult, hidden_size}, ToTensorType(query_scale)); test.AddInput("conv_norm_scale", {hc_mult, hidden_size}, ToTensorType(conv_scale)); - test.AddOutput("output", {1, 1, hc_mult, hidden_size}, ToTensorType(gated_value), false, tolerance, + test.AddOutput("output", key_shape, ToTensorType(gated_value), false, tolerance, tolerance); - test.AddOutput("gated_value_normed", {1, 1, hc_mult, hidden_size}, ToTensorType(expected_normed), false, + test.AddOutput("gated_value_normed", key_shape, ToTensorType(expected_normed), false, tolerance, tolerance); - RunOnSupportedProviders(test); + RunOnSupportedProviders(test, packed); } // Exercises hc_mult > 1 and non-unit norm scales for an arbitrary hidden_size. hidden_size == 4 @@ -225,7 +232,7 @@ void RunEngramGateNormedTest(float tolerance) { // 64 over hidden_size / components), so only a hidden_size above those strides takes a second // iteration and actually accumulates into the per-thread partials. template -void RunEngramGateVectorizedTest(float tolerance, int64_t hidden) { +void RunEngramGateVectorizedTest(float tolerance, int64_t hidden, bool packed = false) { if (!IsTypeSupported()) { GTEST_SKIP() << "No execution provider available for this type"; } @@ -271,14 +278,18 @@ void RunEngramGateVectorizedTest(float tolerance, int64_t hidden) { OpTester test("EngramGate", 1, kMSDomain); test.AddAttribute("epsilon", kEpsilon); - test.AddInput("key", {kBatch, kSequence, kHcMult, hidden}, ToTensorType(key)); - test.AddInput("query", {kBatch, kSequence, kHcMult, hidden}, ToTensorType(query)); - test.AddInput("value", {kBatch, kSequence, hidden}, ToTensorType(value)); + const std::vector key_shape = packed ? std::vector{kBatch * kSequence, kHcMult, hidden} + : std::vector{kBatch, kSequence, kHcMult, hidden}; + const std::vector value_shape = packed ? std::vector{kBatch * kSequence, hidden} + : std::vector{kBatch, kSequence, hidden}; + test.AddInput("key", key_shape, ToTensorType(key)); + test.AddInput("query", key_shape, ToTensorType(query)); + test.AddInput("value", value_shape, ToTensorType(value)); test.AddInput("key_norm_scale", {kHcMult, hidden}, ToTensorType(key_scale)); test.AddInput("query_norm_scale", {kHcMult, hidden}, ToTensorType(query_scale)); - test.AddOutput("output", {kBatch, kSequence, kHcMult, hidden}, ToTensorType(expected), false, tolerance, + test.AddOutput("output", key_shape, ToTensorType(expected), false, tolerance, tolerance); - RunOnSupportedProviders(test); + RunOnSupportedProviders(test, packed); } // --------------------------------------------------------------------------------------------- @@ -1793,6 +1804,14 @@ TEST(EngramOpsTest, EngramGateVectorizedFloat16) { RunEngramGateVectorizedTest(3e-3f, 4); } +TEST(EngramOpsTest, EngramGatePackedFloat) { + RunEngramGateVectorizedTest(1e-4f, 4, true); +} + +TEST(EngramOpsTest, EngramGatePackedFloat16) { + RunEngramGateVectorizedTest(3e-3f, 4, true); +} + // 260 channels is the smallest multiple of 4 above both accumulation strides (CUDA's blockDim.x of // 256 and, with components == 4, the WGSL workgroup size of 64 over hidden_size / 4 == 65), so this // is the only case where either strided reduction loop runs more than once per thread. @@ -1816,6 +1835,10 @@ TEST(EngramOpsTest, EngramGateNormedFloat16) { RunEngramGateNormedTest(2e-3f); } +TEST(EngramOpsTest, EngramGatePackedNormedFloat) { + RunEngramGateNormedTest(1e-4f, true); +} + TEST(EngramOpsTest, EngramGateNormedBFloat16) { RunEngramGateNormedTest(2e-2f); } diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc new file mode 100644 index 0000000000000..2173434746ce9 --- /dev/null +++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc @@ -0,0 +1,283 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#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(FusedTensorOpsTest, ArgMaxLastAxisWithoutKeepdims) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + OpTester test("ArgMaxLastAxis", 1, kMSDomain); + test.AddInput("logits", {2, 4}, {1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f}); + test.AddOutput("indices", {2}, {1, 0}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(FusedTensorOpsTest, ArgMaxLastAxisWithKeepdims) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + OpTester test("ArgMaxLastAxis", 1, kMSDomain); + test.AddAttribute("keepdims", 1); + test.AddInput("logits", {2, 4}, + ToTensorType({1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f})); + test.AddOutput("indices", {2, 1}, {1, 0}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(FusedTensorOpsTest, ArgMaxLastAxisPartitionedVocabulary) { + constexpr int64_t width = 248321; + for (bool keepdims : {false, true}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + std::vector values(3 * width, -std::numeric_limits::infinity()); + values[4095] = values[4096] = values[width - 1] = 5.0f; + values[2 * width - 1] = 7.0f; + OpTester test("ArgMaxLastAxis", 1, kMSDomain); + test.AddAttribute("keepdims", keepdims ? 1 : 0); + test.AddInput("logits", {3, width}, ToBFloat16(values)); + test.AddOutput("indices", keepdims ? std::vector{3, 1} : std::vector{3}, + {4095, width - 1, 0}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } +} + +TEST(FusedTensorOpsTest, EmbeddingRMSNorm) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + const std::vector table = {3.0f, 4.0f, 0.0f, 0.0f, 0.0f, 0.0f, 2.0f, 0.0f}; + const float inverse_first = 1.0f / std::sqrt(6.25f + 1e-6f); + const float inverse_second = 1.0f / std::sqrt(1.0f + 1e-6f); + const std::vector expected = { + 3.0f * inverse_first, 4.0f * inverse_first, 0.0f, 0.0f, + 0.0f, 0.0f, 2.0f * inverse_second, 0.0f}; + + OpTester test("EmbeddingRMSNorm", 1, kMSDomain); + test.AddAttribute("epsilon", 1e-6f); + test.AddInput("ids", {2}, {0, 1}); + test.AddInput("table", {2, 4}, ToTensorType(table)); + test.AddInput("weight", {4}, {1.0f, 1.0f, 1.0f, 1.0f}); + test.AddOutput("normalized", {2, 4}, ToTensorType(expected)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(FusedTensorOpsTest, RowSelect) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + OpTester test("RowSelect", 1, kMSDomain); + test.AddInput("selection", {}, {1}); + for (int i = 0; i < 3; ++i) { + const std::string input_name = "x" + std::to_string(i); + const std::string output_name = "y" + std::to_string(i); + test.AddInput(input_name.c_str(), {1, 2, 2}, + ToTensorType({1.0f, 2.0f, 3.0f, 4.0f})); + test.AddOutput(output_name.c_str(), {1, 1, 2}, + ToTensorType({3.0f, 4.0f})); + } + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(FusedTensorOpsTest, SigmoidGateMul) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute("head_dim", 2); + test.AddAttribute("num_heads", 1); + test.AddAttribute("q_stride", 4); + test.AddAttribute("tok_stride", 4); + test.AddAttribute("q_offset", 0); + test.AddInput("attention", {1, 1, 2}, ToTensorType({2.0f, 4.0f})); + test.AddInput("projection", {1, 4}, ToTensorType({1.0f, 2.0f, 0.0f, 0.0f})); + test.AddOutput("output", {1, 1, 2}, ToTensorType({1.0f, 2.0f})); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +void RunSigmoidGateMulBroadcastResidualRounding(bool sigmoid_in_fp32, bool silu_gate = false) { + for (bool residual : {false, true}) { + for (int64_t width : {1, 257, 2560}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + const std::vector base_gates = ToBFloat16({-100.0f, -5.5f, -0.0f, 0.3f, 5.5f, 100.0f}); + std::vector gates(silu_gate ? 6 * width : 6); + for (size_t i = 0; i < gates.size(); ++i) gates[i] = base_gates[i % 6]; + std::vector values(6 * width), residuals(6 * width), expected(6 * width); + for (size_t i = 0; i < values.size(); ++i) { + values[i] = BFloat16(static_cast(static_cast(i % 23) - 11) * 0.27f); + residuals[i] = BFloat16(static_cast(static_cast(i % 17) - 8) * 0.13f); + const float gate = gates[silu_gate ? i : i / width].ToFloat(); + const float exp_value = std::exp(gate >= 0 ? -gate : gate); + const float sigmoid = gate >= 0 ? 1.0f / (1.0f + exp_value) : exp_value / (1.0f + exp_value); + const BFloat16 denominator(1.0f + BFloat16(exp_value).ToFloat()); + const BFloat16 reciprocal(1.0f / denominator.ToFloat()); + const BFloat16 rounded_sigmoid = sigmoid_in_fp32 + ? BFloat16(sigmoid) + : (gate > 0 ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat())); + const BFloat16 activation = silu_gate ? BFloat16(gate * rounded_sigmoid.ToFloat()) : rounded_sigmoid; + const BFloat16 product(values[i].ToFloat() * activation.ToFloat()); + expected[i] = residual ? BFloat16(product.ToFloat() + residuals[i].ToFloat()) : product; + } + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute(silu_gate ? "silu_gate" : "broadcast_gate", 1); + test.AddAttribute("sigmoid_in_fp32", sigmoid_in_fp32 ? 1 : 0); + test.AddInput("attention", {2, 3, width}, values); + test.AddInput("projection", {2, 3, silu_gate ? width : 1}, gates); + if (residual) test.AddInput("residual", {2, 3, width}, residuals); + test.AddOutput("output", {2, 3, width}, expected); + test.SetOutputTolerance(0.0f, 0.0f); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + } +} + +TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastResidualRounding) { + RunSigmoidGateMulBroadcastResidualRounding(true); +} + +TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastInputPrecision) { + RunSigmoidGateMulBroadcastResidualRounding(false); +} + +TEST(FusedTensorOpsTest, SigmoidGateMulSiluRounding) { + RunSigmoidGateMulBroadcastResidualRounding(true, true); + RunSigmoidGateMulBroadcastResidualRounding(false, true); +} + +TEST(FusedTensorOpsTest, SigmoidGateMulRejectsBroadcastAndResidualShapes) { + for (bool bad_residual : {false, true}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute("broadcast_gate", 1); + test.AddInput("attention", {2, 3}, ToBFloat16(std::vector(6, 1.0f))); + test.AddInput("projection", bad_residual ? std::vector{2, 1} : std::vector{1, 2}, + ToBFloat16({1.0f, 1.0f})); + if (bad_residual) test.AddInput("residual", {1}, ToBFloat16({1.0f})); + test.AddOutput("output", {2, 3}, ToBFloat16(std::vector(6, 0.0f))); + test.Run(OpTester::ExpectResult::kExpectFailure, + bad_residual ? "residual must have the same shape" : "singleton final dimension", + {}, nullptr, &providers); + } +} + +TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastEmptyAndInvalidAttributes) { + for (int mode = 0; mode < 4; ++mode) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute("broadcast_gate", mode == 1 ? 2 : 1); + if (mode == 2) test.AddAttribute("sigmoid_in_fp32", 2); + if (mode == 3) test.AddAttribute("head_dim", 2); + test.AddInput("attention", {0, 3}, {}); + test.AddInput("projection", {0, 1}, {}); + test.AddInput("residual", {0, 3}, {}); + test.AddOutput("output", {0, 3}, {}); + const std::vector errors = {"", "broadcast_gate must be zero or one", + "sigmoid_in_fp32 must be zero or one", + "cannot be combined with packed layout attributes"}; + test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, + errors[mode], {}, nullptr, &providers); + } +} + +TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) { + for (int mode = 0; mode < 5; ++mode) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute("silu_gate", mode == 1 ? 2 : 1); + if (mode == 2) test.AddAttribute("broadcast_gate", 1); + if (mode == 3) test.AddAttribute("head_dim", 2); + test.AddInput("attention", {0, 3}, {}); + test.AddInput("projection", {0, mode == 4 ? 1 : 3}, {}); + test.AddOutput("output", {0, 3}, {}); + const std::vector errors = {"", "silu_gate must be zero or one", "are exclusive", + "cannot be combined with packed layout attributes", + "silu gate must have the same shape"}; + test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, + errors[mode], {}, nullptr, &providers); + } +} + +TEST(FusedTensorOpsTest, UntileHyperConnectionWeights) { + constexpr int64_t kInputElements = 160ll * 20480; + constexpr int64_t kVectorCount = kInputElements / 16; + for (const int64_t kind : {0, 1}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + OpTester test("UntileHyperConnectionWeights", 1, kMSDomain); + test.AddAttribute("kind", kind); + std::vector input(kInputElements); + std::vector expected(kInputElements); + for (int64_t i = 0; i < kInputElements; ++i) { + input[i] = BFloat16(static_cast(i % 127)); + } + for (int64_t vector_index = 0; vector_index < kVectorCount; ++vector_index) { + int64_t destination_vector; + if (kind == 0) { + const int64_t rp = vector_index & 127; + const int64_t t1 = vector_index >> 7; + const int64_t i = t1 % 10; + const int64_t t2 = t1 / 10; + const int64_t j = t2 % 20; + const int64_t s = t2 / 20; + destination_vector = (((s * 10 + i) * 128 + rp) * 20 + j); + } else { + const int64_t c = vector_index & 3; + const int64_t t1 = vector_index >> 2; + const int64_t rp = t1 & 31; + const int64_t t2 = t1 >> 5; + const int64_t i = t2 % 10; + const int64_t hb = t2 / 10; + destination_vector = (((i * 32 + rp) * 4 + c) * 160 + hb); + } + for (int64_t lane = 0; lane < 16; ++lane) { + expected[destination_vector * 16 + lane] = input[vector_index * 16 + lane]; + } + } + test.AddInput("tiled", {160, 20480}, input); + const std::vector output_shape = kind == 0 + ? std::vector{10240, 320} + : std::vector{320, 10240}; + test.AddOutput("untiled", output_shape, expected); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } +} + +} // namespace +} // namespace onnxruntime::test diff --git a/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc new file mode 100644 index 0000000000000..917d936f45468 --- /dev/null +++ b/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc @@ -0,0 +1,285 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include + +#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 +float RoundToType(float value) { + if constexpr (std::is_same_v) { + return MLFloat16(value).ToFloat(); + } else if constexpr (std::is_same_v) { + return BFloat16(value).ToFloat(); + } else { + return value; + } +} + +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; + } +} + +float Sigmoid(float v) { return 1.0f / (1.0f + std::exp(-v)); } + +struct Attrs { + float epsilon = 1e-6f; + float reduction_scale = 0.25f; + float silu_alpha = 0.25f; + float gate_div = 4.0f; + float gate_mul = 2.0f; +}; + +struct Problem { + int64_t rows, C, H, R; + std::vector streams, norm, down, up, inject, prev_out, prev_gate; +}; + +template +Problem MakeProblem(int64_t rows, int64_t C, int64_t H, int64_t R, bool prologue, unsigned seed) { + std::mt19937 gen(seed); + std::normal_distribution normal(0.0f, 1.0f); + const int64_t K = C * H; + auto fill = [&](size_t n, float scale, float shift) { + std::vector v(n); + for (auto& x : v) x = RoundToType(normal(gen) * scale + shift); + return v; + }; + Problem p{rows, C, H, R, {}, {}, {}, {}, {}, {}, {}}; + p.streams = fill(rows * K, 1.0f, 0.0f); + p.norm = fill(K, 0.1f, 1.0f); + p.down = fill(K * R, 1.0f / std::sqrt(static_cast(K)), 0.0f); + p.up = fill(R * K, 2.0f / std::sqrt(static_cast(R)), 0.0f); + p.inject = fill(K * C, 1.0f / std::sqrt(static_cast(K)), 0.0f); + if (prologue) { + p.prev_out = fill(rows * H, 1.0f, 0.0f); + p.prev_gate = fill(rows * C, 0.5f, 1.0f); + } + return p; +} + +struct Expected { + std::vector pre, gate, streams; +}; + +// float32 math with the rounding points listed in the schema documentation. +template +Expected Reference(const Problem& p, const Attrs& a) { + const int64_t K = p.C * p.H; + auto rd = [](float v) { return RoundToType(v); }; + Expected e; + e.streams = p.streams; + const bool prologue = !p.prev_out.empty(); + for (int64_t r = 0; r < p.rows; ++r) { + std::vector hn(K); + for (int64_t c = 0; c < p.C; ++c) { + float sum_sq = 0.0f; + for (int64_t h = 0; h < p.H; ++h) { + float v = p.streams[r * K + c * p.H + h]; + if (prologue) { + v = rd(v + p.prev_gate[r * p.C + c] * p.prev_out[r * p.H + h]); + e.streams[r * K + c * p.H + h] = v; + } + sum_sq += v * v; + } + const float inv = 1.0f / std::sqrt(sum_sq / static_cast(p.H) + a.epsilon); + for (int64_t h = 0; h < p.H; ++h) { + hn[c * p.H + h] = rd(e.streams[r * K + c * p.H + h] * inv * p.norm[c * p.H + h]); + } + } + std::vector hidden(p.R); + for (int64_t j = 0; j < p.R; ++j) { + double acc = 0.0; + for (int64_t k = 0; k < K; ++k) acc += static_cast(hn[k]) * p.down[k * p.R + j]; + const float z = rd(rd(static_cast(acc)) * a.silu_alpha); + hidden[j] = rd(z * rd(Sigmoid(z))); + } + std::vector gates(K); + for (int64_t k = 0; k < K; ++k) { + double acc = 0.0; + for (int64_t j = 0; j < p.R; ++j) acc += static_cast(hidden[j]) * p.up[j * K + k]; + gates[k] = rd(Sigmoid(rd(static_cast(acc)))); + } + for (int64_t h = 0; h < p.H; ++h) { + float sum = 0.0f; + for (int64_t c = 0; c < p.C; ++c) sum += hn[c * p.H + h] * gates[c * p.H + h]; + e.pre.push_back(rd(sum * a.reduction_scale)); + } + for (int64_t c = 0; c < p.C; ++c) { + double acc = 0.0; + for (int64_t k = 0; k < K; ++k) acc += static_cast(hn[k]) * p.inject[k * p.C + c]; + const float t = rd(rd(static_cast(acc)) / a.gate_div); + e.gate.push_back(rd(rd(Sigmoid(t)) * a.gate_mul)); + } + } + return e; +} + +template +void RunTest(int64_t rows, int64_t C, int64_t H, int64_t R, bool prologue, bool request_streams_out, + float tolerance, unsigned seed = 1) { + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) { + GTEST_SKIP() << "HyperConnectionLowRankMix is only implemented for CUDA"; + } + const Attrs a; + const Problem p = MakeProblem(rows, C, H, R, prologue, seed); + const Expected e = Reference(p, a); + const int64_t K = C * H; + + OpTester test("HyperConnectionLowRankMix", 1, kMSDomain); + test.AddAttribute("num_branches", C); + test.AddAttribute("epsilon", a.epsilon); + test.AddAttribute("reduction_scale", a.reduction_scale); + test.AddAttribute("silu_alpha", a.silu_alpha); + test.AddAttribute("gate_div", a.gate_div); + test.AddAttribute("gate_mul", a.gate_mul); + test.AddInput("streams", {rows, K}, ToTensorType(p.streams)); + test.AddInput("norm_scale", {K}, ToTensorType(p.norm)); + test.AddInput("down_weight", {K, R}, ToTensorType(p.down)); + test.AddInput("up_weight", {R, K}, ToTensorType(p.up)); + test.AddInput("inject_weight", {K, C}, ToTensorType(p.inject)); + if (prologue) { + test.AddInput("prev_block_output", {rows, H}, ToTensorType(p.prev_out)); + test.AddInput("prev_post_gate", {rows, C}, ToTensorType(p.prev_gate)); + } else { + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + } + test.AddOutput("pre_output", {rows, H}, ToTensorType(e.pre)); + test.AddOutput("post_gate", {rows, C}, ToTensorType(e.gate)); + if (request_streams_out) { + test.AddOutput("streams_out", {rows, K}, ToTensorType(e.streams)); + } else { + test.AddOptionalOutputEdge(); + } + test.SetOutputTolerance(tolerance, tolerance); + + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +constexpr float kFp32Tol = 2e-4f; +constexpr float kFp16Tol = 8e-3f; +constexpr float kBf16Tol = 6e-2f; + +} // namespace + +TEST(HyperConnectionLowRankMixTest, FusedFloat16WithPrologue) { + RunTest(2, 4, 16, 8, true, true, kFp16Tol); +} + +TEST(HyperConnectionLowRankMixTest, FusedFloat16WithoutPrologue) { + RunTest(1, 4, 64, 16, false, false, kFp16Tol); +} + +TEST(HyperConnectionLowRankMixTest, FusedBFloat16WithPrologueWithoutStreamsOut) { + RunTest(3, 4, 64, 16, true, false, kBf16Tol); +} + +TEST(HyperConnectionLowRankMixTest, FusedBFloat16TwoBranches) { + RunTest(2, 2, 32, 8, true, true, kBf16Tol); +} + +TEST(HyperConnectionLowRankMixTest, FusedMultipleKSlicesAndWeightChunks) { + RunTest(2, 4, 1288, 40, true, true, kBf16Tol); + RunTest(1, 4, 2560, 320, false, true, kFp16Tol); +} + +TEST(HyperConnectionLowRankMixTest, UnfusedManyRows) { + RunTest(9, 4, 16, 8, true, true, kFp16Tol); + RunTest(17, 4, 64, 16, false, true, kBf16Tol); +} + +TEST(HyperConnectionLowRankMixTest, FusedSpecializedRank320) { + for (int rows : {1, 2, 3}) { + RunTest(rows, 4, 1288, 320, true, true, kBf16Tol); + RunTest(rows, 4, 1288, 320, false, false, kFp16Tol); + } +} + +TEST(HyperConnectionLowRankMixTest, FusedDecodeShapeRank320) { + for (int rows : {1, 2}) { + for (bool prologue : {false, true}) { + RunTest(rows, 4, 2560, 320, prologue, true, kBf16Tol, 37); + RunTest(rows, 4, 2560, 320, prologue, true, kFp16Tol, 53); + } + } +} + +TEST(HyperConnectionLowRankMixTest, UnfusedUnsupportedShape) { + RunTest(2, 3, 10, 5, true, true, kFp16Tol); + RunTest(1, 4, 12, 16, false, false, kBf16Tol); +} + +TEST(HyperConnectionLowRankMixTest, Float32) { + RunTest(1, 4, 64, 16, true, true, kFp32Tol); + RunTest(5, 2, 24, 8, false, true, kFp32Tol); +} + +TEST(HyperConnectionLowRankMixTest, PrologueInputsMustBeProvidedTogether) { + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) { + GTEST_SKIP() << "HyperConnectionLowRankMix is only implemented for CUDA"; + } + constexpr int64_t C = 2, H = 8, R = 8, K = C * H; + OpTester test("HyperConnectionLowRankMix", 1, kMSDomain); + test.AddAttribute("num_branches", C); + test.AddInput("streams", {1, K}, std::vector(K, 1.0f)); + test.AddInput("norm_scale", {K}, std::vector(K, 1.0f)); + test.AddInput("down_weight", {K, R}, std::vector(K * R, 0.1f)); + test.AddInput("up_weight", {R, K}, std::vector(R * K, 0.1f)); + test.AddInput("inject_weight", {K, C}, std::vector(K * C, 0.1f)); + test.AddInput("prev_block_output", {1, H}, std::vector(H, 1.0f)); + test.AddOptionalInputEdge(); + test.AddOutput("pre_output", {1, H}, std::vector(H, 0.0f)); + test.AddOutput("post_gate", {1, C}, std::vector(C, 0.0f)); + test.AddOptionalOutputEdge(); + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectFailure, "must be provided together", {}, nullptr, &providers); +} + +TEST(HyperConnectionLowRankMixTest, WeightShapeMismatchFails) { + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) { + GTEST_SKIP() << "HyperConnectionLowRankMix is only implemented for CUDA"; + } + constexpr int64_t C = 2, H = 8, R = 8, K = C * H; + OpTester test("HyperConnectionLowRankMix", 1, kMSDomain); + test.AddAttribute("num_branches", C); + test.AddInput("streams", {1, K}, std::vector(K, 1.0f)); + test.AddInput("norm_scale", {K}, std::vector(K, 1.0f)); + test.AddInput("down_weight", {K, R}, std::vector(K * R, 0.1f)); + test.AddInput("up_weight", {R + 1, K}, std::vector((R + 1) * K, 0.1f)); + test.AddInput("inject_weight", {K, C}, std::vector(K * C, 0.1f)); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddOutput("pre_output", {1, H}, std::vector(H, 0.0f)); + test.AddOutput("post_gate", {1, C}, std::vector(C, 0.0f)); + test.AddOptionalOutputEdge(); + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectFailure, "up_weight must have shape", {}, nullptr, &providers); +} + +} // namespace onnxruntime::test 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); diff --git a/onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py b/onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py new file mode 100644 index 0000000000000..0709465b55125 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py @@ -0,0 +1,311 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Numeric check for com.microsoft.HyperConnectionLowRankMix against a torch reference. + +The reference is the unfused subgraph the operator replaces, + + [HyperConnectionPostMix] -> BranchwiseRMSNorm -> MatMul(down) -> ScaledSiLU -> MatMul(up) -> Sigmoid + -> HyperConnectionPreMix, and MatMul(inject) -> Div -> Sigmoid -> Mul, + +evaluated with float32 math and rounded to the activation type at the points the standard operators round +(see the schema documentation), with a correctly rounded sigmoid. The fused path (few rows, 16-bit types) and +the cuBLAS based path (many rows, float32, shapes the fused kernels do not support) are both covered. +""" + +# Tensor dimensions keep the upper-case names the schema doc gives them. +# ruff: noqa: N806 + +import unittest + +import numpy as np +import torch +from onnx import TensorProto as TP # noqa: N817 +from onnx import helper + +import onnxruntime as ort + +DTYPES = { + TP.FLOAT: torch.float32, + TP.FLOAT16: torch.float16, + TP.BFLOAT16: torch.bfloat16, +} + +# (num_branches, hidden, rank): the real Qwen3.8-Flash-Next block first, then small and irregular shapes. +REAL_SHAPE = (4, 2560, 320) +SMALL_SHAPES = [ + (4, 64, 16), + (2, 32, 8), + (8, 64, 32), + (1, 48, 8), + (4, 1288, 40), # two K slices per branch in the fused path + (4, 4096, 512), # four K slices per branch and two weight chunks in the up kernel + (3, 40, 12), # not supported by the fused kernels (odd branch count, rank % 8 != 0) + (4, 100, 16), # hidden % 8 != 0 +] +ATTRS = dict(epsilon=1e-6, reduction_scale=0.25, silu_alpha=0.25, gate_div=4.0, gate_mul=2.0) + + +def has_cuda(): + return "CUDAExecutionProvider" in ort.get_available_providers() + + +def build_model(io_elem, C, H, R, prologue, attrs): + """Graph with float32 inputs/outputs; Casts move the activations to the type under test.""" + F = TP.FLOAT + K = C * H + node_inputs = ["streams", "norm_scale", "down_weight", "up_weight", "inject_weight"] + graph_inputs = [ + ("streams", ["rows", K]), + ("norm_scale", [K]), + ("down_weight", [K, R]), + ("up_weight", [R, K]), + ("inject_weight", [K, C]), + ] + if prologue: + node_inputs += ["prev_block_output", "prev_post_gate"] + graph_inputs += [("prev_block_output", ["rows", H]), ("prev_post_gate", ["rows", C])] + outputs = ["pre_output", "post_gate", "streams_out"] + + nodes = [helper.make_node("Cast", [n], [n + "_t"], to=io_elem, name=n + "/cast") for n, _ in graph_inputs] + nodes.append( + helper.make_node( + "HyperConnectionLowRankMix", + [n + "_t" for n in node_inputs], + [n + "_t" for n in outputs], + domain="com.microsoft", + num_branches=C, + **attrs, + ) + ) + nodes += [helper.make_node("Cast", [n + "_t"], [n], to=F, name=n + "/castback") for n in outputs] + graph = helper.make_graph( + nodes, + "hc_low_rank_mix", + [helper.make_tensor_value_info(n, F, shape) for n, shape in graph_inputs], + [ + helper.make_tensor_value_info("pre_output", F, ["rows", H]), + helper.make_tensor_value_info("post_gate", F, ["rows", C]), + helper.make_tensor_value_info("streams_out", F, ["rows", K]), + ], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)] + ) + model.ir_version = 10 + return model + + +def make_inputs(dtype, rows, C, H, R, prologue, seed): + g = torch.Generator().manual_seed(seed) + K = C * H + + def rnd(*shape, scale=1.0, shift=0.0): + return (torch.randn(*shape, generator=g) * scale + shift).to(dtype).float() + + inputs = { + "streams": rnd(rows, K), + "norm_scale": rnd(K, scale=0.1, shift=1.0), + "down_weight": rnd(K, R, scale=K**-0.5), + "up_weight": rnd(R, K, scale=2.0 * R**-0.5), + "inject_weight": rnd(K, C, scale=K**-0.5), + } + if prologue: + inputs["prev_block_output"] = rnd(rows, H) + inputs["prev_post_gate"] = (torch.rand(rows, C, generator=g) * 2.0).to(dtype).float() + return inputs + + +def reference(inputs, dtype, C, H, attrs): + def rd(t): + return t.to(dtype).float() + + def sigmoid(t): + return torch.sigmoid(t.double()).float() + + x = inputs["streams"] + rows = x.shape[0] + s = x.view(rows, C, H) + if "prev_block_output" in inputs: + s = rd(s + inputs["prev_post_gate"][:, :, None] * inputs["prev_block_output"][:, None, :]) + inv_rms = torch.rsqrt(s.pow(2).mean(-1, keepdim=True) + attrs["epsilon"]) + hn = rd(s * inv_rms * inputs["norm_scale"].view(C, H)) + hn_flat = hn.view(rows, C * H) + + down = rd(hn_flat @ inputs["down_weight"]) + z = rd(down * attrs["silu_alpha"]) + hidden = rd(z * rd(sigmoid(z))) + up = rd(hidden @ inputs["up_weight"]) + g = rd(sigmoid(up)).view(rows, C, H) + pre = rd(attrs["reduction_scale"] * (hn * g).sum(1)) + inject = rd(hn_flat @ inputs["inject_weight"]) + gate = rd(rd(sigmoid(rd(inject / attrs["gate_div"]))) * attrs["gate_mul"]) + return pre, gate, s.reshape(rows, C * H) + + +class SessionCache: + def __init__(self): + self.sessions = {} + + def get(self, io_elem, C, H, R, prologue, attrs): + key = (io_elem, C, H, R, prologue, tuple(sorted(attrs.items()))) + if key not in self.sessions: + model = build_model(io_elem, C, H, R, prologue, attrs) + so = ort.SessionOptions() + so.log_severity_level = 3 + self.sessions[key] = ort.InferenceSession( + model.SerializeToString(), so, providers=["CUDAExecutionProvider"] + ) + return self.sessions[key] + + +def tolerances(io_elem): + # One ulp of the output type plus slack for the different fp32 accumulation order flipping a rounding. + return {TP.FLOAT: (1e-3, 1e-5), TP.FLOAT16: (4e-3, 2e-3), TP.BFLOAT16: (3e-2, 8e-3)}[io_elem] + + +@unittest.skipUnless(has_cuda(), "HyperConnectionLowRankMix is a CUDA-only kernel") +class TestHyperConnectionLowRankMix(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.cache = SessionCache() + + def _check(self, io_elem, rows, shape, prologue, attrs=ATTRS, seed=0): + C, H, R = shape + dtype = DTYPES[io_elem] + inputs = make_inputs(dtype, rows, C, H, R, prologue, seed) + session = self.cache.get(io_elem, C, H, R, prologue, attrs) + pre, gate, streams_out = session.run(None, {k: v.numpy() for k, v in inputs.items()}) + exp_pre, exp_gate, exp_streams = reference(inputs, dtype, C, H, attrs) + rtol, atol = tolerances(io_elem) + msg = f"dtype={dtype} rows={rows} shape={shape} prologue={prologue}" + np.testing.assert_allclose(pre, exp_pre.numpy(), rtol=rtol, atol=atol, err_msg="pre_output " + msg) + np.testing.assert_allclose(gate, exp_gate.numpy(), rtol=rtol, atol=atol, err_msg="post_gate " + msg) + # The prologue is a single fp32 operation with one rounding: it matches exactly (up to fma contraction). + np.testing.assert_allclose(streams_out, exp_streams.numpy(), rtol=rtol, atol=atol, err_msg="streams_out " + msg) + return pre, gate, streams_out + + def test_real_shape(self): + for io_elem in DTYPES: + for rows in (1, 2, 4, 16, 128): + for prologue in (False, True): + with self.subTest(io_elem=io_elem, rows=rows, prologue=prologue): + self._check(io_elem, rows, REAL_SHAPE, prologue) + + def test_small_shapes(self): + for io_elem in DTYPES: + for shape in SMALL_SHAPES: + for rows in (1, 2, 16): + for prologue in (False, True): + with self.subTest(io_elem=io_elem, shape=shape, rows=rows, prologue=prologue): + self._check(io_elem, rows, shape, prologue, seed=1) + + def test_default_attributes(self): + attrs = dict(epsilon=1e-5, reduction_scale=1.0, silu_alpha=1.0, gate_div=1.0, gate_mul=1.0) + for io_elem in DTYPES: + for rows in (1, 3): + with self.subTest(io_elem=io_elem, rows=rows): + self._check(io_elem, rows, (4, 64, 16), True, attrs=attrs, seed=2) + + def test_other_gate_constants(self): + attrs = dict(epsilon=1e-6, reduction_scale=0.5, silu_alpha=0.75, gate_div=3.0, gate_mul=1.5) + for io_elem in (TP.FLOAT16, TP.BFLOAT16): + for rows in (1, 2): + with self.subTest(io_elem=io_elem, rows=rows): + self._check(io_elem, rows, REAL_SHAPE, True, attrs=attrs, seed=3) + + def test_chain_of_blocks(self): + """Two consecutive calls through the prologue equal the explicit post-mix between them.""" + C, H, R = REAL_SHAPE + for io_elem in (TP.FLOAT16, TP.BFLOAT16): + dtype = DTYPES[io_elem] + with self.subTest(io_elem=io_elem): + inputs = make_inputs(dtype, 2, C, H, R, False, seed=4) + first = self.cache.get(io_elem, C, H, R, False, ATTRS) + second = self.cache.get(io_elem, C, H, R, True, ATTRS) + _, gate, streams = first.run(None, {k: v.numpy() for k, v in inputs.items()}) + # Use the block "output" as an independent tensor (a stand-in for the attention/MLP result). + block_out = torch.randn(2, H, generator=torch.Generator().manual_seed(8)).to(dtype).float().numpy() + chained = dict(inputs) + chained["streams"] = streams + chained["prev_block_output"] = block_out + chained["prev_post_gate"] = gate + out = second.run(None, {k: np.asarray(v) for k, v in chained.items()}) + ref_inputs = {k: torch.from_numpy(np.asarray(v)) for k, v in chained.items()} + exp = reference(ref_inputs, dtype, C, H, ATTRS) + rtol, atol = tolerances(io_elem) + np.testing.assert_allclose(out[0], exp[0].numpy(), rtol=rtol, atol=atol) + np.testing.assert_allclose(out[1], exp[1].numpy(), rtol=rtol, atol=atol) + + def test_fused_rows_are_independent(self): + """The fused kernels use a fixed reduction order: a row gives the same bits at any batch size.""" + C, H, R = REAL_SHAPE + for io_elem in (TP.FLOAT16, TP.BFLOAT16): + dtype = DTYPES[io_elem] + with self.subTest(io_elem=io_elem): + inputs = make_inputs(dtype, 4, C, H, R, True, seed=5) + batch = self.cache.get(io_elem, C, H, R, True, ATTRS).run( + None, {k: v.numpy() for k, v in inputs.items()} + ) + for r in range(4): + row_inputs = { + k: (v[r : r + 1] if k in ("streams", "prev_block_output", "prev_post_gate") else v).numpy() + for k, v in inputs.items() + } + single = self.cache.get(io_elem, C, H, R, True, ATTRS).run(None, row_inputs) + for b, s in zip(batch, single, strict=True): + np.testing.assert_array_equal(b[r : r + 1], s) + + def test_no_prologue_streams_out_is_passthrough(self): + C, H, R = 4, 64, 16 + inputs = make_inputs(torch.bfloat16, 2, C, H, R, False, seed=6) + session = self.cache.get(TP.BFLOAT16, C, H, R, False, ATTRS) + _, _, streams_out = session.run(None, {k: v.numpy() for k, v in inputs.items()}) + np.testing.assert_array_equal(streams_out, inputs["streams"].numpy()) + + def test_cuda_graph_capture_and_replay(self): + """No host synchronization or unmanaged allocation: the op can be captured and replayed.""" + C, H, R = REAL_SHAPE + for io_elem in (TP.FLOAT16, TP.BFLOAT16): + for rows in (1, 2): + with self.subTest(io_elem=io_elem, rows=rows): + dtype = DTYPES[io_elem] + inputs = make_inputs(dtype, rows, C, H, R, True, seed=9) + model = build_model(io_elem, C, H, R, True, ATTRS) + so = ort.SessionOptions() + so.log_severity_level = 3 + graph_session = ort.InferenceSession( + model.SerializeToString(), + so, + providers=[("CUDAExecutionProvider", {"enable_cuda_graph": "1"})], + ) + binding = graph_session.io_binding() + keep = [] + for name, value in inputs.items(): + ortvalue = ort.OrtValue.ortvalue_from_numpy(value.numpy(), "cuda", 0) + keep.append(ortvalue) + binding.bind_ortvalue_input(name, ortvalue) + shapes = {"pre_output": (rows, H), "post_gate": (rows, C), "streams_out": (rows, C * H)} + outputs = {} + for name, shape in shapes.items(): + outputs[name] = ort.OrtValue.ortvalue_from_shape_and_type(shape, np.float32, "cuda", 0) + binding.bind_ortvalue_output(name, outputs[name]) + graph_session.run_with_iobinding(binding) # captures + graph_session.run_with_iobinding(binding) # replays + expected = self.cache.get(io_elem, C, H, R, True, ATTRS).run( + None, {k: v.numpy() for k, v in inputs.items()} + ) + for name, exp in zip(shapes, expected, strict=True): + np.testing.assert_array_equal(outputs[name].numpy(), exp, err_msg=name) + + def test_zero_rows(self): + C, H, R = 4, 64, 16 + inputs = make_inputs(torch.float16, 0, C, H, R, False, seed=7) + session = self.cache.get(TP.FLOAT16, C, H, R, False, ATTRS) + pre, gate, streams_out = session.run(None, {k: v.numpy() for k, v in inputs.items()}) + self.assertEqual(pre.shape, (0, H)) + self.assertEqual(gate.shape, (0, C)) + self.assertEqual(streams_out.shape, (0, C * H)) + + +if __name__ == "__main__": + unittest.main() From f6ff038f5e65a7d58df9c7700c0a99621a00593f Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 19:25:53 +0000 Subject: [PATCH 02/19] Keep normalization coverage in its dedicated PR --- .../contrib_ops/hyper_connection_op_test.cc | 252 ------------------ 1 file changed, 252 deletions(-) diff --git a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc index 622116282406c..f8f590ec32598 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 { @@ -130,255 +127,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); -} - -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 d9b224f50c7a83fa24929c31304d38c8b4593551 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 20:57:10 +0000 Subject: [PATCH 03/19] Fix packed Engram inference and fused tensor boundary handling --- docs/ContribOperators.md | 259 ++++++++++++++++-- docs/OperatorKernels.md | 10 +- .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 67 +---- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 61 +---- .../contrib_ops/cuda/bert/fused_tensor_ops.h | 14 - .../hyper_connection_low_rank_mix_impl.cu | 8 +- .../core/graph/contrib_ops/bert_defs.cc | 33 ++- .../test/contrib_ops/engram_ops_test.cc | 22 ++ .../test/contrib_ops/fused_tensor_ops_test.cc | 60 +++- .../hyper_connection_low_rank_mix_op_test.cc | 62 +++++ 10 files changed, 444 insertions(+), 152 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 1015b41e18389..239e55aafc9e6 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -3,6 +3,8 @@ Do not modify directly.* * com.microsoft + * com.microsoft.ArgMaxLastAxis + * com.microsoft.EmbeddingRMSNorm * com.microsoft.HyperConnectionLowRankMix * com.microsoft.Attention * com.microsoft.AttnLSTM @@ -119,9 +121,11 @@ Do not modify directly.* * com.microsoft.RestorePadding * com.microsoft.Rfft * com.microsoft.RotaryEmbedding + * com.microsoft.RowSelect * com.microsoft.SampleOp * com.microsoft.Sampling * com.microsoft.ScaledSiLU + * com.microsoft.SigmoidGateMul * com.microsoft.SkipGroupNorm * com.microsoft.SkipLayerNormalization * com.microsoft.SkipSimplifiedLayerNormalization @@ -136,6 +140,7 @@ Do not modify directly.* * com.microsoft.Trilu * com.microsoft.UnfoldTensor * com.microsoft.Unique + * com.microsoft.UntileHyperConnectionWeights * com.microsoft.VarlenCausalConvWithState * com.microsoft.VarlenNGramHashMapping * com.microsoft.WhisperBeamSearch @@ -144,6 +149,45 @@ Do not modify directly.* * experimental com.microsoft.QEmbedLayerNormalization ## com.microsoft +### **com.microsoft.ArgMaxLastAxis** + + CUDA contrib operators for fused tensor operations used by transformer models. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
keepdims : int (default is 0)
+
Keep a singleton final axis when nonzero.
+
+ +#### Inputs + +
+
logits : T
+
Logits with vocabulary on the last axis.
+
+ +#### Outputs + +
+
indices : I
+
First maximum index for each leading row. NaNs are ignored; an all-NaN row returns index zero.
+
+ +#### Type Constraints + +
+
T : tensor(float), tensor(float16), tensor(bfloat16)
+
Supported logit types.
+
I : tensor(int64)
+
Argmax indices are int64.
+
+ + ### **com.microsoft.Attention** Multi-Head Attention that can be either unidirectional (like GPT-2) or bidirectional (like BERT). @@ -1971,6 +2015,51 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.EmbeddingRMSNorm** + + CUDA contrib operators for fused tensor operations used by transformer models. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
epsilon : float
+
RMS normalization epsilon.
+
+ +#### Inputs + +
+
ids : I
+
Embedding row indices with any shape. Negative indices in [-vocabulary, -1] count from the end. Out-of-range indices produce NaN output rows.
+
table : T
+
BF16 embedding matrix [vocabulary, hidden].
+
weight : F
+
FP32 RMS normalization weight [hidden].
+
+ +#### Outputs + +
+
normalized : T
+
BF16 normalized embeddings [shape(ids), hidden].
+
+ +#### Type Constraints + +
+
I : tensor(int64)
+
Indices use int64.
+
T : tensor(bfloat16)
+
Embedding and result use BF16.
+
F : tensor(float)
+
Normalization weight uses FP32.
+
+ + ### **com.microsoft.EngramGate** Fuses the Engram gate. @@ -1999,7 +2088,7 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
epsilon : float
+
epsilon : float (default is 1e-05)
Epsilon used by both RMS normalization steps. Default is 1e-5.
@@ -2007,11 +2096,11 @@ This version of the operator has been available since version 1 of the 'com.micr
key : T
-
Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size).
+
Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size) or (total_tokens, hc_mult, hidden_size).
query : T
-
Hidden-state queries with shape (batch_size, sequence_length, hc_mult, hidden_size).
+
Hidden-state queries with the same shape as key.
value : T
-
Projected Engram value shared by every hyper-connection, with shape (batch_size, sequence_length, hidden_size).
+
Projected Engram value shared by every hyper-connection, with shape (batch_size, sequence_length, hidden_size) or (total_tokens, hidden_size).
key_norm_scale : T
RMSNorm scale for keys with shape (hc_mult, hidden_size).
query_norm_scale : T
@@ -2024,9 +2113,9 @@ This version of the operator has been available since version 1 of the 'com.micr
output : T
-
Gated value tensor with shape (batch_size, sequence_length, hc_mult, hidden_size).
+
Gated value tensor with the same shape as key.
gated_value_normed (optional) : T
-
Optional RMS-normalized gated value tensor with shape (batch_size, sequence_length, hc_mult, hidden_size).
+
Optional RMS-normalized gated value tensor with shape (batch_size, sequence_length, hc_mult, hidden_size) or (total_tokens, hc_mult, hidden_size).
#### Type Constraints @@ -7224,6 +7313,48 @@ This version of the operator has been available since version 1 of the 'com.micr
+### **com.microsoft.RowSelect** + + CUDA contrib operators for fused tensor operations used by transformer models. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Inputs + +
+
selection : I
+
Scalar-like int64 row selector; zero selects row 0.
+
x0 : T
+
BF16 tensor with shape [1, 2, ...].
+
x1 : T
+
BF16 tensor with shape [1, 2, ...].
+
x2 : T
+
BF16 tensor with shape [1, 2, ...].
+
+ +#### Outputs + +
+
y0 : T
+
Selected BF16 tensor with shape [1, 1, ...].
+
y1 : T
+
Selected BF16 tensor with shape [1, 1, ...].
+
y2 : T
+
Selected BF16 tensor with shape [1, 1, ...].
+
+ +#### Type Constraints + +
+
I : tensor(int64)
+
Selector uses int64.
+
T : tensor(bfloat16)
+
Inputs and outputs use BF16.
+
+ + ### **com.microsoft.SampleOp** Sample echo operator. @@ -7383,6 +7514,61 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.SigmoidGateMul** + + CUDA contrib operators for fused tensor operations used by transformer models. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
broadcast_gate : int (default is 0)
+
When one, projection has the attention shape with final dimension one. Packed layout attributes must be zero. CUDA only.
+
head_dim : int (default is 0)
+
Dimension of each attention head.
+
num_heads : int (default is 0)
+
Number of attention heads.
+
q_offset : int (default is 0)
+
Offset of the query projection within each token.
+
q_stride : int (default is 0)
+
Per-head stride in the packed projection.
+
sigmoid_in_fp32 : int (default is 1)
+
One computes sigmoid in FP32 before BF16 conversion. Zero preserves the input-type intermediate arithmetic of CUDA Sigmoid, including subtraction on negative inputs.
+
silu_gate : int (default is 0)
+
When one, projection and attention have identical shapes. Multiply the BF16-rounded sigmoid by projection and round to BF16 before multiplying attention. Cannot be combined with broadcast_gate or packed layout attributes.
+
tok_stride : int (default is 0)
+
Per-token stride in the packed projection.
+
+ +#### Inputs (2 - 3) + +
+
attention : T
+
BF16 attention values.
+
projection : T
+
BF16 packed projection containing gate values.
+
residual (optional) : T
+
Optional BF16 residual with the attention shape.
+
+ +#### Outputs + +
+
output : T
+
Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 before optionally adding residual and rounding again.
+
+ +#### Type Constraints + +
+
T : tensor(bfloat16)
+
Inputs and output use BF16.
+
+ + ### **com.microsoft.SkipGroupNorm** This operator element-wise adds x, skip and bias, then apply group normalization and optional activation. @@ -8288,6 +8474,45 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.UntileHyperConnectionWeights** + + CUDA contrib operators for fused tensor operations used by transformer models. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
kind : int
+
Untile layout kind, either 0 or 1.
+
+ +#### Inputs (1 - 2) + +
+
tiled : T
+
Tiled BF16 weight [160, 20480].
+
unused (optional) : T
+
Reserved optional input retained for exporter compatibility.
+
+ +#### Outputs + +
+
untiled : T
+
Untiled BF16 weight.
+
+ +#### Type Constraints + +
+
T : tensor(bfloat16)
+
Input and output use BF16.
+
+ + ### **com.microsoft.VarlenCausalConvWithState** Stateful causal depthwise convolution over a packed, token-major batch of variable-length @@ -8778,20 +9003,20 @@ No versioning maintained for experimental ops. One hyper-connection block with a low-rank gate network, fused into one operator. C = num_branches streams of width H are kept as one flattened (..., C * H) tensor. With the optional prologue the operator first applies the post-mix of the previous block (identity stream mix): - + S[..., c, h] = streams[..., c, h] + prev_post_gate[..., c] * prev_block_output[..., h] - + and otherwise S = streams. It then computes, with `rms` taken per branch over H: - + hn[..., c, h] = S[..., c, h] / sqrt(mean_h(S[..., c, :]^2) + epsilon) * norm_scale[c, h] d = ScaledSiLU(hn @ down_weight, silu_alpha) (..., R) g[..., c, h] = sigmoid(d @ up_weight) per stream and feature pre_output = reduction_scale * sum_c(hn[..., c, h] * g[..., c, h]) (..., H) post_gate = sigmoid((hn @ inject_weight) / gate_div) * gate_mul (..., C) - + pre_output is the input of the next transformer block and post_gate is prev_post_gate of the next call. streams_out is S (the streams after the prologue; a copy of streams when there is no prologue). - + Numerics contract. Matrix products and reductions accumulate in float32. The result of every step is rounded to T once, as the output of the equivalent standard operator would be: S = cast_T(streams + prev_post_gate * prev_block_output), hn = cast_T(S * inv_rms * norm_scale), @@ -8809,17 +9034,17 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
epsilon : float
+
epsilon : float (default is 1e-05)
Epsilon added to the mean square of each branch before the reciprocal square root.
-
gate_div : float
+
gate_div : float (default is 1.0)
Divisor applied to the inject logits before the sigmoid. Must not be zero.
-
gate_mul : float
+
gate_mul : float (default is 1.0)
Multiplier applied to the sigmoid of the inject logits.
num_branches : int (required)
Number of residual streams C. The last dimension of streams is C * H.
-
reduction_scale : float
+
reduction_scale : float (default is 1.0)
Multiplier applied to the sum over branches of pre_output.
-
silu_alpha : float
+
silu_alpha : float (default is 1.0)
Scale of the ScaledSiLU between the down and up projections.
@@ -8859,3 +9084,5 @@ This version of the operator has been available since version 1 of the 'com.micr
T : tensor(float), tensor(float16), tensor(bfloat16)
Constrain all floating-point inputs and outputs to the same type.
+ + diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 669e3083f6b15..7a0fe2a406ed1 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -586,7 +586,7 @@ The **OpSet Version** column uses the following notation: |DynamicQuantizeMatMul|*in* A:**T1**
*in* B:**T2**
*in* b_scale:**T1**
*in* b_zero_point:**T2**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(float)
**T2** = tensor(int8), tensor(uint8)| |DynamicTimeWarping|*in* input:**F**
*out* output:**I**|1+|**F** = tensor(float)
**I** = tensor(int32)| |EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float)| -|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**T** = tensor(float), tensor(float16)| +|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**MLFloat16** = tensor(float16)
**float** = tensor(float)| |ExpandDims|*in* X:**T**
*in* axis:**tensor(int32)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**axis** = tensor(int32)| |FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)| @@ -1079,6 +1079,7 @@ The **OpSet Version** column uses the following notation: | | | | |**Operator Domain:** *com.microsoft*|||| +|ArgMaxLastAxis|*in* logits:**T**
*out* indices:**I**|1+|**BFloat16** = tensor(bfloat16)
**I** = tensor(int64)
**MLFloat16** = tensor(float16)
**float** = tensor(float)| |Attention|*in* input:**T**
*in* weights:**T**
*in* bias:**T**
*in* mask_index:**M**
*in* past:**T**
*in* attention_bias:**T**
*in* past_sequence_length:**M**
*out* output:**T**
*out* present:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |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), tensor(float16)| |BiasAdd|*in* X:**T**
*in* bias:**T**
*in* skip:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| @@ -1101,7 +1102,8 @@ The **OpSet Version** column uses the following notation: |DynamicSparseAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* auxiliary_key:**T**
*in* auxiliary_value:**T**
*in* selected_indices:**M**
*in* selected_counts:**M**
*in* seqlens_k:**M**
*in* total_sequence_length:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* position_ids:**tensor(int64)**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* head_sink:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |DynamicTimeWarping|*in* input:**F**
*out* output:**I**|1+|**F** = tensor(float)
**I** = tensor(int32)| |EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)| -|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|EmbeddingRMSNorm|*in* ids:**I**
*in* table:**T**
*in* weight:**F**
*out* normalized:**T**|1+|**F** = tensor(float)
**I** = tensor(int64)
**T** = tensor(bfloat16)| +|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**BFloat16** = tensor(bfloat16)
**MLFloat16** = tensor(float16)
**float** = tensor(float)| |FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| @@ -1117,6 +1119,7 @@ The **OpSet Version** column uses the following notation: |GridSample|*in* X:**T1**
*in* Grid:**T1**
*out* Y:**T2**|1+|**T1** = tensor(float)
**T2** = tensor(float)| |GroupNorm|*in* X:**T**
*in* gamma:**M**
*in* beta:**M**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |GroupQueryAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T_CACHE**
*in* past_value:**T_CACHE**
*in* seqlens_k:**M**
*in* total_sequence_length:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* position_ids:**tensor(int64)**
*in* attention_bias:**T**
*in* head_sink:**T**
*in* k_scale:**T_KV_SCALE**
*in* v_scale:**T_KV_SCALE**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*out* output:**T**
*out* present_key:**T_CACHE**
*out* present_value:**T_CACHE**
*out* output_qk:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)
**T_CACHE** = tensor(bfloat16), tensor(float16), tensor(float8e4m3fn), tensor(int8), tensor(uint8)
**T_KV_SCALE** = tensor(float)| +|HyperConnectionLowRankMix|*in* streams:**T**
*in* norm_scale:**T**
*in* down_weight:**T**
*in* up_weight:**T**
*in* inject_weight:**T**
*in* prev_block_output:**T**
*in* prev_post_gate:**T**
*out* pre_output:**T**
*out* post_gate:**T**
*out* streams_out:**T**|1+|**BFloat16** = tensor(bfloat16)
**MLFloat16** = tensor(float16)
**float** = tensor(float)| |HyperConnectionPostMix|*in* streams:**T**
*in* block_output:**T**
*in* post_mix:**M**
*in* stream_mix:**M**
*out* output:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |HyperConnectionPreMix|*in* streams:**T**
*in* pre_mix:**M**
*out* output:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |Inverse|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| @@ -1154,8 +1157,10 @@ The **OpSet Version** column uses the following notation: |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)| +|RowSelect|*in* selection:**I**
*in* x0:**T**
*in* x1:**T**
*in* x2:**T**
*out* y0:**T**
*out* y1:**T**
*out* y2:**T**|1+|**I** = tensor(int64)
**T** = tensor(bfloat16)| |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)| +|SigmoidGateMul|*in* attention:**T**
*in* projection:**T**
*in* residual:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16)| |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)| |SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |SkipSimplifiedLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| @@ -1165,6 +1170,7 @@ The **OpSet Version** column uses the following notation: |TransposeMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |Trilu|*in* X:**T**
*in* k:**tensor(int64)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |UnfoldTensor|*in* input:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|UntileHyperConnectionWeights|*in* tiled:**T**
*in* unused:**T**
*out* untiled:**T**|1+|**T** = tensor(bfloat16)| |VarlenCausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* cumulative_sequence_length:**M**
*in* bias:**T**
*in* initial_state:**T**
*in* capture_count:**M**
*out* output:**T**
*out* final_state:**T**
*out* state_update:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |VarlenNGramHashMapping|*in* input_ids:**M**
*in* multipliers:**M**
*in* vocab_sizes:**M**
*in* cumulative_sequence_length:**S**
*in* past_ids:**M**
*in* head_offsets:**M**
*in* eos_token_id:**M**
*in* segment_ids:**tensor(int32)**
*in* past_segment_ids:**S**
*in* capture_count:**S**
*out* hash_ids:**M**
*out* present_ids:**M**
*out* present_segment_ids:**S**
*out* state_update:**M**|1+|**M** = tensor(int32), tensor(int64)
**S** = tensor(int32)| |WhisperBeamSearch|*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**
*in* cross_qk_layer_head:**I**
*in* extra_decoding_ids:**I**
*in* temperature:**T**
*out* sequences:**I**
*out* sequences_scores:**T**
*out* scores:**T**
*out* cross_qk:**V**
*out* non_speech_probs:**T**|1+|**T** = tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc index 555c98f1fb16d..6248f17a92905 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -7,7 +7,7 @@ #include #include -#include "core/common/float8.h" +#include "core/common/safeint.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" @@ -48,16 +48,6 @@ ONNX_OPERATOR_KERNEL_EX( .TypeConstraint("T", DataTypeImpl::GetTensorType()), SigmoidGateMul); -#if !defined(DISABLE_FLOAT8_TYPES) -ONNX_OPERATOR_TYPED_KERNEL_EX( - DequantizeLinear, kMSDomain, 1, Float8E4M3FN_BFloat16, kCudaExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) - .TypeConstraint("TS", DataTypeImpl::GetTensorType()), - DequantizeLinearFp8Bf16); -#endif - ONNX_OPERATOR_KERNEL_EX( UntileHyperConnectionWeights, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) @@ -153,7 +143,7 @@ SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) const int64_t silu = info.GetAttrOrDefault("silu_gate", 0); ORT_ENFORCE(silu == 0 || silu == 1, "SigmoidGateMul silu_gate must be zero or one"); ORT_ENFORCE(!(broadcast_gate_ && silu_gate_), "SigmoidGateMul broadcast_gate and silu_gate are exclusive"); - ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && q_stride_ >= 2 * head_dim_ && tok_stride_ > 0 && q_offset_ >= 0), + ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && q_stride_ >= 2LL * head_dim_ && tok_stride_ > 0 && q_offset_ >= 0), "SigmoidGateMul received invalid layout attributes"); ORT_ENFORCE(!(broadcast_gate_ || silu_gate_) || (head_dim_ == 0 && heads_ == 0 && q_stride_ == 0 && tok_stride_ == 0 && q_offset_ == 0), @@ -184,10 +174,17 @@ Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(count % row_width == 0, "SigmoidGateMul attention size must be divisible by heads * head_dim"); const int64_t tokens = count / row_width; - ORT_RETURN_IF_NOT(projection->Shape().Size() >= - (tokens - 1) * tok_stride_ + q_offset_ + - static_cast(heads_) * q_stride_, - "SigmoidGateMul projection is too small for the configured layout"); + if (tokens > 0) { + int64_t required_extent; + int64_t head_extent; + ORT_RETURN_IF_NOT(SafeMultiply(tokens - 1, tok_stride_, required_extent) && + SafeMultiply(static_cast(heads_), static_cast(q_stride_), head_extent) && + SafeAdd(required_extent, q_offset_, required_extent) && + SafeAdd(required_extent, head_extent, required_extent), + "SigmoidGateMul projection extent overflows int64"); + ORT_RETURN_IF_NOT(projection->Shape().Size() >= required_extent, + "SigmoidGateMul projection is too small for the configured layout"); + } } Tensor* output = context->Output(0, attention->Shape()); return LaunchSigmoidGateMul( @@ -199,44 +196,6 @@ Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { sigmoid_in_fp32_, silu_gate_); } -DequantizeLinearFp8Bf16::DequantizeLinearFp8Bf16(const OpKernelInfo& info) - : CudaKernel(info), axis_(info.GetAttrOrDefault("axis", 1)) {} - -Status DequantizeLinearFp8Bf16::ComputeInternal(OpKernelContext* context) const { - const Tensor* input = context->Input(0); - const Tensor* scales = context->Input(1); - const Tensor* zero_point = context->Input(2); - const auto& input_shape = input->Shape(); - ORT_RETURN_IF_NOT(input_shape.NumDimensions() > 0, - "DequantizeLinear FP8/BF16 input must have rank at least one"); - int64_t axis = axis_ < 0 ? axis_ + static_cast(input_shape.NumDimensions()) : axis_; - - int64_t axis_size = 1; - int64_t inner_size = 1; - if (scales->Shape().Size() != 1) { - ORT_RETURN_IF_NOT(axis >= 0 && axis < static_cast(input_shape.NumDimensions()), - "DequantizeLinear axis is out of range"); - ORT_RETURN_IF_NOT(scales->Shape().NumDimensions() == 1 && - scales->Shape().Size() == input_shape[axis], - "DequantizeLinear scale must be scalar or match the selected axis"); - axis_size = scales->Shape().Size(); - inner_size = input_shape.SizeFromDimension(static_cast(axis + 1)); - } - if (zero_point != nullptr) { - ORT_RETURN_IF_NOT(zero_point->Shape().Size() == 1 || - zero_point->Shape() == scales->Shape(), - "DequantizeLinear zero point must be scalar or match the scale"); - } - - Tensor* output = context->Output(0, input_shape); - return LaunchDequantizeLinearFp8Bf16( - Stream(context), reinterpret_cast(input->DataRaw()), - zero_point == nullptr ? nullptr : reinterpret_cast(zero_point->DataRaw()), - scales->Data(), - reinterpret_cast(output->MutableData()), input_shape.Size(), - axis_size, inner_size, zero_point == nullptr || zero_point->Shape().Size() == 1); -} - UntileHyperConnectionWeights::UntileHyperConnectionWeights(const OpKernelInfo& info) : CudaKernel(info), kind_(static_cast(info.GetAttrOrDefault("kind", -1))) { ORT_ENFORCE(kind_ == 0 || kind_ == 1, "UntileHyperConnectionWeights kind must be 0 or 1"); diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu index 789edc4a26615..8878f442185b3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -58,7 +58,7 @@ __global__ void ArgmaxRowsKernel(const T* input, const int64_t* input_indices, __syncthreads(); } if (tid == 0) { - output[blockIdx.x] = indices[0]; + output[blockIdx.x] = partial_values == nullptr && indices[0] == INT64_MAX ? 0 : indices[0]; if (partial_values != nullptr) partial_values[blockIdx.x] = values[0]; } } @@ -84,10 +84,17 @@ Status LaunchTypedArgmax(cudaStream_t stream, const T* input, int64_t rows, __global__ void EmbedNormKernel(const int64_t* ids, const uint16_t* table, const float* weight, float epsilon, - uint16_t* output, int64_t rows, int64_t hidden_size) { + uint16_t* output, int64_t vocab_size, int64_t hidden_size) { __shared__ float sums[512]; const int64_t row_index = blockIdx.x; - const int64_t token_id = ids[row_index]; + int64_t token_id = ids[row_index]; + if (token_id < -vocab_size || token_id >= vocab_size) { + for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { + output[row_index * hidden_size + col] = __bfloat16_as_ushort(__float2bfloat16_rn(CUDART_NAN_F)); + } + return; + } + if (token_id < 0) token_id += vocab_size; const uint16_t* embedding = table + token_id * hidden_size; float sum = 0.0f; for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { @@ -109,7 +116,6 @@ __global__ void EmbedNormKernel(const int64_t* ids, const uint16_t* table, output[row_index * hidden_size + col] = __bfloat16_as_ushort(__float2bfloat16_rn(normalized)); } - (void)rows; } __global__ void RowSelectKernel(const int64_t* selection, @@ -184,37 +190,6 @@ __global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention, output[i] = value; } -__device__ __forceinline__ float DecodeE4M3FN(uint8_t value) { - const float sign = (value & 0x80u) != 0 ? -1.0f : 1.0f; - const int exponent = (value >> 3) & 0x0f; - const int mantissa = value & 0x07; - if (exponent == 0) { - return sign * ldexpf(static_cast(mantissa), -9); - } - if (exponent == 0x0f && mantissa == 0x07) { - return CUDART_NAN_F; - } - return sign * ldexpf(1.0f + static_cast(mantissa) * 0.125f, - exponent - 7); -} - -__global__ void DequantizeLinearFp8Bf16Kernel(const uint8_t* input, - const uint8_t* zero_point, - const float* scales, uint16_t* output, - int64_t count, int64_t axis_size, - int64_t inner_size, - bool scalar_zero_point) { - const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (index >= count) return; - float value = DecodeE4M3FN(input[index]); - const int64_t scale_index = axis_size == 1 ? 0 : (index / inner_size) % axis_size; - if (zero_point != nullptr) { - value -= DecodeE4M3FN(zero_point[scalar_zero_point ? 0 : scale_index]); - } - value *= scales[scale_index]; - output[index] = __bfloat16_as_ushort(__float2bfloat16_rn(value)); -} - __global__ void UntileHcBf16Kernel(const uint16_t* input, uint16_t* output, int kind) { constexpr int64_t kVectorCount = 204800; constexpr int64_t kElementCount = kVectorCount * 16; @@ -275,7 +250,7 @@ Status LaunchEmbeddingRMSNorm(cudaStream_t stream, const int64_t* ids, const uin "EmbeddingRMSNorm has too many rows for a CUDA grid"); if (rows == 0) return Status::OK(); EmbedNormKernel<<(rows), 512, 0, stream>>>( - ids, table, weight, epsilon, output, rows, hidden_size); + ids, table, weight, epsilon, output, vocab_size, hidden_size); return CUDA_CALL(cudaGetLastError()); } @@ -326,20 +301,6 @@ Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, return CUDA_CALL(cudaGetLastError()); } -Status LaunchDequantizeLinearFp8Bf16(cudaStream_t stream, const uint8_t* input, - const uint8_t* zero_point, const float* scales, - uint16_t* output, int64_t count, - int64_t axis_size, int64_t inner_size, - bool scalar_zero_point) { - if (count == 0) return Status::OK(); - const int64_t block_count = (count - 1) / kThreads + 1; - ORT_RETURN_IF_NOT(block_count <= std::numeric_limits::max(), - "DequantizeLinear FP8/BF16 requires too many CUDA blocks"); - DequantizeLinearFp8Bf16Kernel<<(block_count), kThreads, 0, stream>>>( - input, zero_point, scales, output, count, axis_size, inner_size, scalar_zero_point); - return CUDA_CALL(cudaGetLastError()); -} - Status LaunchUntileHyperConnectionWeights(cudaStream_t stream, const uint16_t* input, uint16_t* output, int kind) { constexpr int64_t kElementCount = 204800ll * 16; diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h index 7cc6b05922159..46bb3fb756677 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h @@ -27,11 +27,6 @@ Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, int64_t token_stride, int64_t q_offset, int64_t gate_row_width, const uint16_t* residual, bool sigmoid_in_fp32, bool silu_gate); -Status LaunchDequantizeLinearFp8Bf16(cudaStream_t stream, const uint8_t* input, - const uint8_t* zero_point, const float* scales, - uint16_t* output, int64_t count, - int64_t axis_size, int64_t inner_size, - bool scalar_zero_point); Status LaunchUntileHyperConnectionWeights(cudaStream_t stream, const uint16_t* input, uint16_t* output, int kind); @@ -76,15 +71,6 @@ class SigmoidGateMul final : public onnxruntime::cuda::CudaKernel { bool silu_gate_; }; -class DequantizeLinearFp8Bf16 final : public onnxruntime::cuda::CudaKernel { - public: - explicit DequantizeLinearFp8Bf16(const OpKernelInfo& info); - Status ComputeInternal(OpKernelContext* context) const override; - - private: - int64_t axis_; -}; - class UntileHyperConnectionWeights final : public onnxruntime::cuda::CudaKernel { public: explicit UntileHyperConnectionWeights(const OpKernelInfo& info); diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu index 25f9ffccc4194..7e6e2197f66fe 100644 --- a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu @@ -11,10 +11,7 @@ // HcUpPreKernel grid (ceil(H / 16), 1, row groups) // reduces the partials in a fixed order, then ScaledSiLU, up GEMV, sigmoid, pre-mix and the // inject gate. -// Both kernels are launched with programmatic dependent launch (sm_90+, plain launches elsewhere). Each one -// triggers its dependents immediately and prefetches its weights (which never depend on a previous kernel) -// into shared memory before griddepcontrol.wait; everything that consumes another kernel's output comes after -// the wait, so correctness never depends on the trigger position. +// Both kernels use ordinary stream-ordered launches: their weights may be outputs of preceding nodes. // Every intermediate is rounded to T at the points where the unfused ONNX subgraph rounds; the GEMVs // accumulate in fp32 in a fixed (deterministic) order. // @@ -820,7 +817,8 @@ Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, c g.splits, g.partial_stride, config.epsilon}; UpParams up{config.rows, config.branches, config.hidden, config.rank, g.splits, g.partial_stride, config.silu_alpha, config.reduction_scale, config.gate_div, config.gate_mul}; - const bool use_pdl = device_prop.major >= 9; + // Weights can be produced by preceding nodes, so no input prefetch may overlap their execution. + const bool use_pdl = false; if (config.rows == 1) { return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); } diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 5e9751c7d445a..cf0e49d4d15fb 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4588,7 +4588,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Input(2, "value", "Projected Engram value shared by every hyper-connection, with shape " - "(batch_size, sequence_length, hidden_size).", + "(batch_size, sequence_length, hidden_size) or (total_tokens, hidden_size).", "T") .Input(3, "key_norm_scale", @@ -4611,7 +4611,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Output(1, "gated_value_normed", "Optional RMS-normalized gated value tensor with shape " - "(batch_size, sequence_length, hc_mult, hidden_size).", + "(batch_size, sequence_length, hc_mult, hidden_size) or (total_tokens, hc_mult, hidden_size).", "T", OpSchema::Optional) .TypeConstraint("T", @@ -4625,8 +4625,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( if (hasInputShape(ctx, 0)) { const auto& key_shape = getInputShape(ctx, 0); - if (key_shape.dim_size() != 4) { - fail_shape_inference("EngramGate: key must have rank 4"); + if (key_shape.dim_size() != 3 && key_shape.dim_size() != 4) { + fail_shape_inference("EngramGate: key must have rank 3 or 4"); } propagateShapeFromInputToOutput(ctx, 0, 0); if (ctx.getNumOutputs() > 1) { @@ -4635,14 +4635,21 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } if (hasInputShape(ctx, 1)) { const auto& query_shape = getInputShape(ctx, 1); - if (query_shape.dim_size() != 4) { - fail_shape_inference("EngramGate: query must have rank 4"); + if (query_shape.dim_size() != 3 && query_shape.dim_size() != 4) { + fail_shape_inference("EngramGate: query must have rank 3 or 4"); + } + if (hasInputShape(ctx, 0) && query_shape.dim_size() != getInputShape(ctx, 0).dim_size()) { + fail_shape_inference("EngramGate: query rank must match key rank"); } } if (hasInputShape(ctx, 2)) { const auto& value_shape = getInputShape(ctx, 2); - if (value_shape.dim_size() != 3) { - fail_shape_inference("EngramGate: value must have rank 3"); + if (value_shape.dim_size() != 2 && value_shape.dim_size() != 3) { + fail_shape_inference("EngramGate: value must have rank 2 or 3"); + } + if ((hasInputShape(ctx, 0) && value_shape.dim_size() + 1 != getInputShape(ctx, 0).dim_size()) || + (hasInputShape(ctx, 1) && value_shape.dim_size() + 1 != getInputShape(ctx, 1).dim_size())) { + fail_shape_inference("EngramGate: value rank must be one lower than key and query rank"); } } })); @@ -5636,7 +5643,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("keepdims", "Keep a singleton final axis when nonzero.", AttributeProto::INT, static_cast(0)) .Input(0, "logits", "Logits with vocabulary on the last axis.", "T") - .Output(0, "indices", "First maximum index for each leading row.", "I") + .Output(0, "indices", + "First maximum index for each leading row. NaNs are ignored; " + "an all-NaN row returns index zero.", + "I") .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Supported logit types.") .TypeConstraint("I", {"tensor(int64)"}, "Argmax indices are int64.") @@ -5662,7 +5672,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema() .SetDoc(FusedTensorOps_ver1_doc) .Attr("epsilon", "RMS normalization epsilon.", AttributeProto::FLOAT) - .Input(0, "ids", "Embedding row indices with any shape.", "I") + .Input(0, "ids", + "Embedding row indices with any shape. Negative indices in [-vocabulary, -1] " + "count from the end. Out-of-range indices produce NaN output rows.", + "I") .Input(1, "table", "BF16 embedding matrix [vocabulary, hidden].", "T") .Input(2, "weight", "FP32 RMS normalization weight [hidden].", "F") .Output(0, "normalized", "BF16 normalized embeddings [shape(ids), hidden].", "T") diff --git a/onnxruntime/test/contrib_ops/engram_ops_test.cc b/onnxruntime/test/contrib_ops/engram_ops_test.cc index 6f64d835d7e9e..a6c868bbf11ff 100644 --- a/onnxruntime/test/contrib_ops/engram_ops_test.cc +++ b/onnxruntime/test/contrib_ops/engram_ops_test.cc @@ -1812,6 +1812,24 @@ TEST(EngramOpsTest, EngramGatePackedFloat16) { RunEngramGateVectorizedTest(3e-3f, 4, true); } +TEST(EngramOpsTest, EngramGatePackedBFloat16) { + RunEngramGateVectorizedTest(2e-2f, 4, true); +} + +TEST(EngramOpsTest, EngramGateRejectsMixedPackedRanks) { + for (bool bad_query : {false, true}) { + OpTester test("EngramGate", 1, kMSDomain); + test.AddInput("key", {1, 1, 2}, {1, 1}); + test.AddInput("query", bad_query ? std::vector{1, 1, 1, 2} : std::vector{1, 1, 2}, {1, 1}); + test.AddInput("value", bad_query ? std::vector{1, 2} : std::vector{1, 1, 2}, {1, 1}); + test.AddInput("key_norm_scale", {1, 2}, {1, 1}); + test.AddInput("query_norm_scale", {1, 2}, {1, 1}); + test.AddOutput("output", {1, 1, 2}, {0, 0}); + test.Run(OpTester::ExpectResult::kExpectFailure, + bad_query ? "query rank must match key rank" : "value rank must be one lower"); + } +} + // 260 channels is the smallest multiple of 4 above both accumulation strides (CUDA's blockDim.x of // 256 and, with components == 4, the WGSL workgroup size of 64 over hidden_size / 4 == 65), so this // is the only case where either strided reduction loop runs more than once per thread. @@ -1839,6 +1857,10 @@ TEST(EngramOpsTest, EngramGatePackedNormedFloat) { RunEngramGateNormedTest(1e-4f, true); } +TEST(EngramOpsTest, EngramGatePackedNormedBFloat16) { + RunEngramGateNormedTest(2e-2f, true); +} + TEST(EngramOpsTest, EngramGateNormedBFloat16) { RunEngramGateNormedTest(2e-2f); } diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc index 2173434746ce9..accc3f2dea38c 100644 --- a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc +++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc @@ -8,7 +8,6 @@ #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" @@ -75,6 +74,29 @@ TEST(FusedTensorOpsTest, ArgMaxLastAxisPartitionedVocabulary) { } } +template +void RunArgMaxNaNRows() { + for (int64_t width : {1, 4096, 4097, 12289}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + const float nan = std::numeric_limits::quiet_NaN(); + std::vector values(4 * width, nan); + values[2 * width - 1] = -std::numeric_limits::infinity(); + values[2 * width] = values[3 * width - 1] = 5.0f; + values[4 * width - 1] = 7.0f; + OpTester test("ArgMaxLastAxis", 1, kMSDomain); + test.AddInput("logits", {4, width}, ToTensorType(values)); + test.AddOutput("indices", {4}, {0, width - 1, 0, width - 1}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } +} + +TEST(FusedTensorOpsTest, ArgMaxLastAxisNaNSingleAndMultiplePartitions) { + RunArgMaxNaNRows(); + RunArgMaxNaNRows(); + RunArgMaxNaNRows(); +} + TEST(FusedTensorOpsTest, EmbeddingRMSNorm) { auto providers = CudaProvider(); if (providers.empty()) GTEST_SKIP(); @@ -95,6 +117,22 @@ TEST(FusedTensorOpsTest, EmbeddingRMSNorm) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } +TEST(FusedTensorOpsTest, EmbeddingRMSNormIndexBounds) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + const float nan = std::numeric_limits::quiet_NaN(); + OpTester test("EmbeddingRMSNorm", 1, kMSDomain); + test.AddAttribute("epsilon", 0.0f); + test.AddInput("ids", {8}, {0, 1, -1, -2, 2, -3, std::numeric_limits::min(), std::numeric_limits::max()}); + test.AddInput("table", {2, 4}, ToBFloat16({2, 0, 0, 0, 0, 2, 0, 0})); + test.AddInput("weight", {4}, {1, 1, 1, 1}); + test.AddOutput("normalized", {8, 4}, + ToBFloat16({2, 0, 0, 0, 0, 2, 0, 0, 0, 2, 0, 0, 2, 0, 0, 0, + nan, nan, nan, nan, nan, nan, nan, nan, + nan, nan, nan, nan, nan, nan, nan, nan})); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + TEST(FusedTensorOpsTest, RowSelect) { auto providers = CudaProvider(); if (providers.empty()) GTEST_SKIP(); @@ -128,6 +166,26 @@ TEST(FusedTensorOpsTest, SigmoidGateMul) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } +TEST(FusedTensorOpsTest, SigmoidGateMulCheckedProjectionExtent) { + for (int mode = 0; mode < 4; ++mode) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + const bool empty = mode == 3; + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute("head_dim", 1); + test.AddAttribute("num_heads", 1); + test.AddAttribute("q_stride", 2); + test.AddAttribute("tok_stride", mode == 1 ? 1 : std::numeric_limits::max()); + test.AddAttribute("q_offset", mode == 1 ? std::numeric_limits::max() : 0); + const int64_t tokens = empty ? 0 : (mode == 2 ? 3 : 2); + test.AddInput("attention", {tokens, 1}, std::vector(tokens)); + test.AddInput("projection", {0}, {}); + test.AddOutput("output", {tokens, 1}, std::vector(tokens)); + test.Run(empty ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, + empty ? "" : "projection extent overflows int64", {}, nullptr, &providers); + } +} + void RunSigmoidGateMulBroadcastResidualRounding(bool sigmoid_in_fp32, bool silu_gate = false) { for (bool residual : {false, true}) { for (int64_t width : {1, 257, 2560}) { diff --git a/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc index 917d936f45468..65cf62aea7ebf 100644 --- a/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc +++ b/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc @@ -8,6 +8,8 @@ #include #include +#include "core/framework/session_options.h" +#include "core/graph/graph.h" #include "gtest/gtest.h" #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" @@ -182,8 +184,68 @@ constexpr float kFp32Tol = 2e-4f; constexpr float kFp16Tol = 8e-3f; constexpr float kBf16Tol = 6e-2f; +class DependentNormScaleTester : public OpTester { + public: + DependentNormScaleTester() : OpTester("HyperConnectionLowRankMix", 1, kMSDomain) {} + + void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, + std::vector>&) override { + auto* norm = &graph.GetOrCreateNodeArg("produced_norm", nullptr); + auto* producer_gate = &graph.GetOrCreateNodeArg("producer_gate", nullptr); + const Attrs a; + auto add_mix = [&](const char* name, const std::vector& node_inputs, + const std::vector& node_outputs) { + auto& node = graph.AddNode(name, "HyperConnectionLowRankMix", "", node_inputs, node_outputs, nullptr, kMSDomain); + node.AddAttribute("num_branches", int64_t{1}); + node.AddAttribute("epsilon", a.epsilon); + node.AddAttribute("reduction_scale", a.reduction_scale); + node.AddAttribute("silu_alpha", a.silu_alpha); + node.AddAttribute("gate_div", a.gate_div); + node.AddAttribute("gate_mul", a.gate_mul); + }; + add_mix("producer", {inputs[0], inputs[1], inputs[2], inputs[3], inputs[4]}, {norm, producer_gate}); + add_mix("consumer", {inputs[5], norm, inputs[6], inputs[7], inputs[8]}, outputs); + graph.SetOutputs(outputs); + } +}; + +template +void RunDependentNormScaleTest(float tolerance) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP(); + for (unsigned seed : {17u, 29u, 53u}) { + const Attrs a; + const Problem producer = MakeProblem(1, 1, 64, 8, false, seed); + Problem consumer = MakeProblem(1, 1, 64, 8, false, seed + 1); + consumer.norm = Reference(producer, a).pre; + const Expected expected = Reference(consumer, a); + DependentNormScaleTester test; + test.AddInput("producer_streams", {64}, ToTensorType(producer.streams)); + test.AddInput("producer_norm", {64}, ToTensorType(producer.norm)); + test.AddInput("producer_down", {64, 8}, ToTensorType(producer.down)); + test.AddInput("producer_up", {8, 64}, ToTensorType(producer.up)); + test.AddInput("producer_inject", {64, 1}, ToTensorType(producer.inject)); + test.AddInput("consumer_streams", {64}, ToTensorType(consumer.streams)); + test.AddInput("consumer_down", {64, 8}, ToTensorType(consumer.down)); + test.AddInput("consumer_up", {8, 64}, ToTensorType(consumer.up)); + test.AddInput("consumer_inject", {64, 1}, ToTensorType(consumer.inject)); + test.AddOutput("pre_output", {64}, ToTensorType(expected.pre)); + test.AddOutput("post_gate", {1}, ToTensorType(expected.gate)); + test.SetOutputTolerance(tolerance, tolerance); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } +} + } // namespace +TEST(HyperConnectionLowRankMixTest, FusedNormScaleProducedByPrecedingMix) { + RunDependentNormScaleTest(kFp16Tol); + RunDependentNormScaleTest(kBf16Tol); +} + TEST(HyperConnectionLowRankMixTest, FusedFloat16WithPrologue) { RunTest(2, 4, 16, 8, true, true, kFp16Tol); } From e587c56161de193f06fc722e890b8cb144c39c52 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 22:06:21 +0000 Subject: [PATCH 04/19] Fix fused gate bounds, WebGPU fallback, and generated operator docs --- docs/ContribOperators.md | 180 +++++++++--------- docs/OperatorKernels.md | 8 +- .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 10 +- .../webgpu/webgpu_execution_provider.cc | 8 + .../test/contrib_ops/engram_ops_test.cc | 78 ++++++++ .../test/contrib_ops/fused_tensor_ops_test.cc | 51 +++++ 6 files changed, 237 insertions(+), 98 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 239e55aafc9e6..2db268555a57b 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4,8 +4,6 @@ Do not modify directly.* * com.microsoft * com.microsoft.ArgMaxLastAxis - * com.microsoft.EmbeddingRMSNorm - * com.microsoft.HyperConnectionLowRankMix * com.microsoft.Attention * com.microsoft.AttnLSTM * com.microsoft.BeamSearch @@ -36,6 +34,7 @@ Do not modify directly.* * com.microsoft.DynamicTimeWarping * com.microsoft.EPContext * com.microsoft.EmbedLayerNormalization + * com.microsoft.EmbeddingRMSNorm * com.microsoft.EngramGate * com.microsoft.ExpandDims * com.microsoft.FastGelu @@ -58,6 +57,7 @@ Do not modify directly.* * com.microsoft.GridSample * com.microsoft.GroupNorm * com.microsoft.GroupQueryAttention + * com.microsoft.HyperConnectionLowRankMix * com.microsoft.HyperConnectionPostMix * com.microsoft.HyperConnectionPreMix * com.microsoft.Inverse @@ -3423,6 +3423,94 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.HyperConnectionLowRankMix** + + One hyper-connection block with a low-rank gate network, fused into one operator. C = num_branches streams of + width H are kept as one flattened (..., C * H) tensor. With the optional prologue the operator first applies the + post-mix of the previous block (identity stream mix): + + S[..., c, h] = streams[..., c, h] + prev_post_gate[..., c] * prev_block_output[..., h] + + and otherwise S = streams. It then computes, with `rms` taken per branch over H: + + hn[..., c, h] = S[..., c, h] / sqrt(mean_h(S[..., c, :]^2) + epsilon) * norm_scale[c, h] + d = ScaledSiLU(hn @ down_weight, silu_alpha) (..., R) + g[..., c, h] = sigmoid(d @ up_weight) per stream and feature + pre_output = reduction_scale * sum_c(hn[..., c, h] * g[..., c, h]) (..., H) + post_gate = sigmoid((hn @ inject_weight) / gate_div) * gate_mul (..., C) + + pre_output is the input of the next transformer block and post_gate is prev_post_gate of the next call. + streams_out is S (the streams after the prologue; a copy of streams when there is no prologue). + + Numerics contract. Matrix products and reductions accumulate in float32. The result of every step is rounded + to T once, as the output of the equivalent standard operator would be: + S = cast_T(streams + prev_post_gate * prev_block_output), hn = cast_T(S * inv_rms * norm_scale), + down = cast_T(hn @ down_weight), Z = cast_T(silu_alpha * down), SG = cast_T(sigmoid(Z)), hidden = cast_T(Z * SG), + up = cast_T(hidden @ up_weight), g = cast_T(sigmoid(up)), pre_output = cast_T(reduction_scale * sum_c(hn * g)), + inject = cast_T(hn @ inject_weight), post_gate = cast_T(cast_T(sigmoid(cast_T(inject / gate_div))) * gate_mul). + sigmoid is evaluated in float32 with expf and rounded once to T; it does not reproduce a Sigmoid kernel that + rounds to T after every elementary operation. The accumulation order of the matrix products is fixed, so the + results are deterministic, but they are not bitwise identical to cuBLAS. + +#### Version + +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 added to the mean square of each branch before the reciprocal square root.
+
gate_div : float (default is 1.0)
+
Divisor applied to the inject logits before the sigmoid. Must not be zero.
+
gate_mul : float (default is 1.0)
+
Multiplier applied to the sigmoid of the inject logits.
+
num_branches : int (required)
+
Number of residual streams C. The last dimension of streams is C * H.
+
reduction_scale : float (default is 1.0)
+
Multiplier applied to the sum over branches of pre_output.
+
silu_alpha : float (default is 1.0)
+
Scale of the ScaledSiLU between the down and up projections.
+
+ +#### Inputs (5 - 7) + +
+
streams : T
+
Flattened streams (..., C * H). With the prologue these are the streams before the previous block's post-mix.
+
norm_scale : T
+
Per-branch RMS norm weight with shape (C * H).
+
down_weight : T
+
Down projection with shape (C * H, R), row-major [K, N].
+
up_weight : T
+
Up projection with shape (R, C * H), row-major [K, N].
+
inject_weight : T
+
Post-gate projection with shape (C * H, C), row-major [K, N].
+
prev_block_output (optional) : T
+
Optional output of the previous block with shape (..., H). Provide together with prev_post_gate.
+
prev_post_gate (optional) : T
+
Optional post-mix gate of the previous block with shape (..., C). Provide together with prev_block_output.
+
+ +#### Outputs (2 - 3) + +
+
pre_output : T
+
Block input with shape (..., H).
+
post_gate : T
+
Post-mix gate for the next call with shape (..., C).
+
streams_out (optional) : T
+
Optional streams after the prologue, with the same shape as streams. Equal to streams when the prologue inputs are absent.
+
+ +#### Type Constraints + +
+
T : tensor(float), tensor(float16), tensor(bfloat16)
+
Constrain all floating-point inputs and outputs to the same type.
+
+ + ### **com.microsoft.HyperConnectionPostMix** Mixes existing streams and injects one branch output: @@ -8998,91 +9086,3 @@ No versioning maintained for experimental ops.
Constrain input and output types to float32 tensors.
-### **com.microsoft.HyperConnectionLowRankMix** - - One hyper-connection block with a low-rank gate network, fused into one operator. C = num_branches streams of - width H are kept as one flattened (..., C * H) tensor. With the optional prologue the operator first applies the - post-mix of the previous block (identity stream mix): - - S[..., c, h] = streams[..., c, h] + prev_post_gate[..., c] * prev_block_output[..., h] - - and otherwise S = streams. It then computes, with `rms` taken per branch over H: - - hn[..., c, h] = S[..., c, h] / sqrt(mean_h(S[..., c, :]^2) + epsilon) * norm_scale[c, h] - d = ScaledSiLU(hn @ down_weight, silu_alpha) (..., R) - g[..., c, h] = sigmoid(d @ up_weight) per stream and feature - pre_output = reduction_scale * sum_c(hn[..., c, h] * g[..., c, h]) (..., H) - post_gate = sigmoid((hn @ inject_weight) / gate_div) * gate_mul (..., C) - - pre_output is the input of the next transformer block and post_gate is prev_post_gate of the next call. - streams_out is S (the streams after the prologue; a copy of streams when there is no prologue). - - Numerics contract. Matrix products and reductions accumulate in float32. The result of every step is rounded - to T once, as the output of the equivalent standard operator would be: - S = cast_T(streams + prev_post_gate * prev_block_output), hn = cast_T(S * inv_rms * norm_scale), - down = cast_T(hn @ down_weight), Z = cast_T(silu_alpha * down), SG = cast_T(sigmoid(Z)), hidden = cast_T(Z * SG), - up = cast_T(hidden @ up_weight), g = cast_T(sigmoid(up)), pre_output = cast_T(reduction_scale * sum_c(hn * g)), - inject = cast_T(hn @ inject_weight), post_gate = cast_T(cast_T(sigmoid(cast_T(inject / gate_div))) * gate_mul). - sigmoid is evaluated in float32 with expf and rounded once to T; it does not reproduce a Sigmoid kernel that - rounds to T after every elementary operation. The accumulation order of the matrix products is fixed, so the - results are deterministic, but they are not bitwise identical to cuBLAS. - -#### Version - -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 added to the mean square of each branch before the reciprocal square root.
-
gate_div : float (default is 1.0)
-
Divisor applied to the inject logits before the sigmoid. Must not be zero.
-
gate_mul : float (default is 1.0)
-
Multiplier applied to the sigmoid of the inject logits.
-
num_branches : int (required)
-
Number of residual streams C. The last dimension of streams is C * H.
-
reduction_scale : float (default is 1.0)
-
Multiplier applied to the sum over branches of pre_output.
-
silu_alpha : float (default is 1.0)
-
Scale of the ScaledSiLU between the down and up projections.
-
- -#### Inputs (5 - 7) - -
-
streams : T
-
Flattened streams (..., C * H). With the prologue these are the streams before the previous block's post-mix.
-
norm_scale : T
-
Per-branch RMS norm weight with shape (C * H).
-
down_weight : T
-
Down projection with shape (C * H, R), row-major [K, N].
-
up_weight : T
-
Up projection with shape (R, C * H), row-major [K, N].
-
inject_weight : T
-
Post-gate projection with shape (C * H, C), row-major [K, N].
-
prev_block_output (optional) : T
-
Optional output of the previous block with shape (..., H). Provide together with prev_post_gate.
-
prev_post_gate (optional) : T
-
Optional post-mix gate of the previous block with shape (..., C). Provide together with prev_block_output.
-
- -#### Outputs (2 - 3) - -
-
pre_output : T
-
Block input with shape (..., H).
-
post_gate : T
-
Post-mix gate for the next call with shape (..., C).
-
streams_out (optional) : T
-
Optional streams after the prologue, with the same shape as streams. Equal to streams when the prologue inputs are absent.
-
- -#### Type Constraints - -
-
T : tensor(float), tensor(float16), tensor(bfloat16)
-
Constrain all floating-point inputs and outputs to the same type.
-
- - diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 7a0fe2a406ed1..67f6bc2b4558d 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -586,7 +586,7 @@ The **OpSet Version** column uses the following notation: |DynamicQuantizeMatMul|*in* A:**T1**
*in* B:**T2**
*in* b_scale:**T1**
*in* b_zero_point:**T2**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(float)
**T2** = tensor(int8), tensor(uint8)| |DynamicTimeWarping|*in* input:**F**
*out* output:**I**|1+|**F** = tensor(float)
**I** = tensor(int32)| |EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float)| -|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**MLFloat16** = tensor(float16)
**float** = tensor(float)| +|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**T** = tensor(float), tensor(float16)| |ExpandDims|*in* X:**T**
*in* axis:**tensor(int32)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**axis** = tensor(int32)| |FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)| @@ -1079,7 +1079,7 @@ The **OpSet Version** column uses the following notation: | | | | |**Operator Domain:** *com.microsoft*|||| -|ArgMaxLastAxis|*in* logits:**T**
*out* indices:**I**|1+|**BFloat16** = tensor(bfloat16)
**I** = tensor(int64)
**MLFloat16** = tensor(float16)
**float** = tensor(float)| +|ArgMaxLastAxis|*in* logits:**T**
*out* indices:**I**|1+|**I** = tensor(int64)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |Attention|*in* input:**T**
*in* weights:**T**
*in* bias:**T**
*in* mask_index:**M**
*in* past:**T**
*in* attention_bias:**T**
*in* past_sequence_length:**M**
*out* output:**T**
*out* present:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |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), tensor(float16)| |BiasAdd|*in* X:**T**
*in* bias:**T**
*in* skip:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| @@ -1103,7 +1103,7 @@ The **OpSet Version** column uses the following notation: |DynamicTimeWarping|*in* input:**F**
*out* output:**I**|1+|**F** = tensor(float)
**I** = tensor(int32)| |EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)| |EmbeddingRMSNorm|*in* ids:**I**
*in* table:**T**
*in* weight:**F**
*out* normalized:**T**|1+|**F** = tensor(float)
**I** = tensor(int64)
**T** = tensor(bfloat16)| -|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**BFloat16** = tensor(bfloat16)
**MLFloat16** = tensor(float16)
**float** = tensor(float)| +|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| @@ -1119,7 +1119,7 @@ The **OpSet Version** column uses the following notation: |GridSample|*in* X:**T1**
*in* Grid:**T1**
*out* Y:**T2**|1+|**T1** = tensor(float)
**T2** = tensor(float)| |GroupNorm|*in* X:**T**
*in* gamma:**M**
*in* beta:**M**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |GroupQueryAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T_CACHE**
*in* past_value:**T_CACHE**
*in* seqlens_k:**M**
*in* total_sequence_length:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* position_ids:**tensor(int64)**
*in* attention_bias:**T**
*in* head_sink:**T**
*in* k_scale:**T_KV_SCALE**
*in* v_scale:**T_KV_SCALE**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*out* output:**T**
*out* present_key:**T_CACHE**
*out* present_value:**T_CACHE**
*out* output_qk:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)
**T_CACHE** = tensor(bfloat16), tensor(float16), tensor(float8e4m3fn), tensor(int8), tensor(uint8)
**T_KV_SCALE** = tensor(float)| -|HyperConnectionLowRankMix|*in* streams:**T**
*in* norm_scale:**T**
*in* down_weight:**T**
*in* up_weight:**T**
*in* inject_weight:**T**
*in* prev_block_output:**T**
*in* prev_post_gate:**T**
*out* pre_output:**T**
*out* post_gate:**T**
*out* streams_out:**T**|1+|**BFloat16** = tensor(bfloat16)
**MLFloat16** = tensor(float16)
**float** = tensor(float)| +|HyperConnectionLowRankMix|*in* streams:**T**
*in* norm_scale:**T**
*in* down_weight:**T**
*in* up_weight:**T**
*in* inject_weight:**T**
*in* prev_block_output:**T**
*in* prev_post_gate:**T**
*out* pre_output:**T**
*out* post_gate:**T**
*out* streams_out:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |HyperConnectionPostMix|*in* streams:**T**
*in* block_output:**T**
*in* post_mix:**M**
*in* stream_mix:**M**
*out* output:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |HyperConnectionPreMix|*in* streams:**T**
*in* pre_mix:**M**
*out* output:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |Inverse|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc index 6248f17a92905..f2cc3aac368f0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -7,6 +7,7 @@ #include #include +#include "core/common/narrow.h" #include "core/common/safeint.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" @@ -128,9 +129,9 @@ Status RowSelect::ComputeInternal(OpKernelContext* context) const { SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) : CudaKernel(info), - head_dim_(static_cast(info.GetAttrOrDefault("head_dim", 0))), - heads_(static_cast(info.GetAttrOrDefault("num_heads", 0))), - q_stride_(static_cast(info.GetAttrOrDefault("q_stride", 0))), + head_dim_(narrow(info.GetAttrOrDefault("head_dim", 0))), + heads_(narrow(info.GetAttrOrDefault("num_heads", 0))), + q_stride_(narrow(info.GetAttrOrDefault("q_stride", 0))), tok_stride_(info.GetAttrOrDefault("tok_stride", 0)), q_offset_(info.GetAttrOrDefault("q_offset", 0)), broadcast_gate_(info.GetAttrOrDefault("broadcast_gate", 0) == 1), @@ -178,7 +179,8 @@ Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { int64_t required_extent; int64_t head_extent; ORT_RETURN_IF_NOT(SafeMultiply(tokens - 1, tok_stride_, required_extent) && - SafeMultiply(static_cast(heads_), static_cast(q_stride_), head_extent) && + SafeMultiply(static_cast(heads_ - 1), static_cast(q_stride_), head_extent) && + SafeAdd(head_extent, 2LL * head_dim_, head_extent) && SafeAdd(required_extent, q_offset_, required_extent) && SafeAdd(required_extent, head_extent, required_extent), "SigmoidGateMul projection extent overflows int64"); diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 493b5eec03c6c..a2a8c65a539d5 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -777,6 +777,14 @@ std::vector> WebGpuExecutionProvider::GetCapa } } + if (node.OpType() == "EngramGate" && node.Domain() == kMSDomain) { + const auto* key_shape = node.InputDefs()[0]->Shape(); + if (key_shape == nullptr || key_shape->dim_size() != 4) { + LOGS(*GetLogger(), INFO) << "EngramGate requires rank-4 key input for WebGPU; falling back to CPU"; + continue; + } + } + // Check for MatMulBnb4 if (node.OpType() == "MatMulBnb4" && node.Domain() == kMSDomain) { // Current implementation only supports the forward case (transB=1). transB defaults to 1. diff --git a/onnxruntime/test/contrib_ops/engram_ops_test.cc b/onnxruntime/test/contrib_ops/engram_ops_test.cc index a6c868bbf11ff..66d3c0a8085f1 100644 --- a/onnxruntime/test/contrib_ops/engram_ops_test.cc +++ b/onnxruntime/test/contrib_ops/engram_ops_test.cc @@ -76,6 +76,7 @@ void RunOnSupportedProviders(OpTester& test, bool packed = false) { providers.push_back(DefaultCudaExecutionProvider()); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } else { + // OpTester preassigns nodes without GetCapability; the session regression below tests fallback. test.Run(OpTester::ExpectResult::kExpectSuccess, "", packed ? std::unordered_set{kWebGpuExecutionProvider} : std::unordered_set{}); @@ -1816,7 +1817,84 @@ TEST(EngramOpsTest, EngramGatePackedBFloat16) { RunEngramGateVectorizedTest(2e-2f, 4, true); } +#ifdef USE_WEBGPU +template +void RunEngramGatePackedWebGpuCpuFallbackTest() { + for (bool normed : {false, true}) { + auto webgpu = DefaultWebGpuExecutionProvider(); + if (webgpu == nullptr) GTEST_SKIP() << "WebGPU execution provider is not available"; + std::unordered_map domain_to_version{{kOnnxDomain, 17}, {kMSDomain, 1}}; + Model model("EngramGatePackedFallback", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, {}, + DefaultLoggingManager().DefaultLogger()); + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + auto* key = builder.MakeInput({2, 2, 1}, ToTensorType({1, 1, 1, 1})); + auto* query = builder.MakeInput({2, 2, 1}, ToTensorType({1, 1, 1, 1})); + auto* value = builder.MakeInput({2, 1}, ToTensorType({2, 4})); + auto* key_scale = builder.MakeInput({2, 1}, ToTensorType({1, 1})); + auto* query_scale = builder.MakeInput({2, 1}, ToTensorType({1, 1})); + std::vector inputs{key, query, value, key_scale, query_scale}; + std::vector outputs{builder.MakeOutput()}; + if (normed) { + inputs.push_back(builder.MakeInput({2, 1}, ToTensorType({1, 1}))); + outputs.push_back(builder.MakeOutput()); + } + auto& node = builder.AddNode("EngramGate", inputs, outputs, kMSDomain); + node.AddAttribute("epsilon", 0.0f); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + ONNX_NAMESPACE::ModelProto model_proto = model.ToProto(); + std::string model_data; + ASSERT_TRUE(model_proto.SerializeToString(&model_data)); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + InferenceSessionWrapper session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(webgpu))); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultCpuExecutionProvider())); + std::istringstream model_stream(model_data); + ASSERT_STATUS_OK(session.Load(model_stream)); + ASSERT_STATUS_OK(session.Initialize()); + size_t gate_nodes = 0; + for (const auto& assigned_node : session.GetSessionState().GetGraphViewer().Nodes()) { + if (assigned_node.OpType() == "EngramGate") { + ++gate_nodes; + EXPECT_EQ(assigned_node.GetExecutionProviderType(), kCpuExecutionProvider); + } + } + ASSERT_EQ(gate_nodes, 1u); + std::vector fetches; + ASSERT_STATUS_OK(session.Run(RunOptions{}, builder.feeds_, builder.output_names_, &fetches)); + ASSERT_EQ(fetches.size(), normed ? 2u : 1u); + const float gate = Sigmoid(1.0f); + const auto expected = ToTensorType({2 * gate, 2 * gate, 4 * gate, 4 * gate}); + const auto actual = fetches[0].Get().DataAsSpan(); + ASSERT_EQ(actual.size(), expected.size()); + for (size_t i = 0; i < actual.size(); ++i) { + EXPECT_NEAR(static_cast(actual[i]), static_cast(expected[i]), 2e-3f); + } + if (normed) { + ASSERT_EQ(fetches[1].Get().Shape(), TensorShape({2, 2, 1})); + for (T element : fetches[1].Get().DataAsSpan()) { + EXPECT_NEAR(static_cast(element), 1.0f, 2e-3f); + } + } + } +} + +TEST(EngramOpsTest, EngramGatePackedWebGpuCpuFallbackFloat) { + RunEngramGatePackedWebGpuCpuFallbackTest(); +} + +TEST(EngramOpsTest, EngramGatePackedWebGpuCpuFallbackFloat16) { + RunEngramGatePackedWebGpuCpuFallbackTest(); +} +#endif + TEST(EngramOpsTest, EngramGateRejectsMixedPackedRanks) { +#if defined(ORT_NO_EXCEPTIONS) + GTEST_SKIP() << "Static shape inference failures abort when exceptions are disabled."; +#endif for (bool bad_query : {false, true}) { OpTester test("EngramGate", 1, kMSDomain); test.AddInput("key", {1, 1, 2}, {1, 1}); diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc index accc3f2dea38c..8b38aced87697 100644 --- a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc +++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc @@ -186,6 +186,51 @@ TEST(FusedTensorOpsTest, SigmoidGateMulCheckedProjectionExtent) { } } +TEST(FusedTensorOpsTest, SigmoidGateMulPaddedHeadStride) { + for (int64_t tokens : {1, 2}) { + for (bool truncated : {false, true}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute("head_dim", 1); + test.AddAttribute("num_heads", 2); + test.AddAttribute("q_stride", 4); + test.AddAttribute("tok_stride", 6); + const int64_t extent = tokens * 6 - (truncated ? 1 : 0); + test.AddInput("attention", {tokens, 2}, + tokens == 1 ? ToBFloat16({2, 4}) : ToBFloat16({2, 4, 6, 8})); + test.AddInput("projection", {extent}, std::vector(extent)); + test.AddOutput("output", {tokens, 2}, + tokens == 1 ? ToBFloat16({1, 2}) : ToBFloat16({1, 2, 3, 4})); + test.Run(truncated ? OpTester::ExpectResult::kExpectFailure : OpTester::ExpectResult::kExpectSuccess, + truncated ? "projection is too small" : "", {}, nullptr, &providers); + } + } +} + +TEST(FusedTensorOpsTest, SigmoidGateMulRejectsOutOfRangeLayoutAttributes) { +#if defined(ORT_NO_EXCEPTIONS) + GTEST_SKIP() << "Checked narrowing failures abort when exceptions are disabled."; +#endif + for (const char* attribute : {"head_dim", "num_heads", "q_stride"}) { + for (int64_t value : {int64_t{4294967300}, int64_t{-4294967292}}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + SCOPED_TRACE(attribute); + SCOPED_TRACE(value); + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute("head_dim", std::string(attribute) == "head_dim" ? value : 2); + test.AddAttribute("num_heads", std::string(attribute) == "num_heads" ? value : 2); + test.AddAttribute("q_stride", std::string(attribute) == "q_stride" ? value : 4); + test.AddAttribute("tok_stride", 8); + test.AddInput("attention", {1, 4}, std::vector(4)); + test.AddInput("projection", {8}, std::vector(8)); + test.AddOutput("output", {1, 4}, std::vector(4)); + test.Run(OpTester::ExpectResult::kExpectFailure, "narrowing_error", {}, nullptr, &providers); + } + } +} + void RunSigmoidGateMulBroadcastResidualRounding(bool sigmoid_in_fp32, bool silu_gate = false) { for (bool residual : {false, true}) { for (int64_t width : {1, 257, 2560}) { @@ -255,6 +300,9 @@ TEST(FusedTensorOpsTest, SigmoidGateMulRejectsBroadcastAndResidualShapes) { TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastEmptyAndInvalidAttributes) { for (int mode = 0; mode < 4; ++mode) { +#if defined(ORT_NO_EXCEPTIONS) + if (mode != 0) continue; +#endif auto providers = CudaProvider(); if (providers.empty()) GTEST_SKIP(); OpTester test("SigmoidGateMul", 1, kMSDomain); @@ -275,6 +323,9 @@ TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastEmptyAndInvalidAttributes) { TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) { for (int mode = 0; mode < 5; ++mode) { +#if defined(ORT_NO_EXCEPTIONS) + if (mode >= 1 && mode <= 3) continue; +#endif auto providers = CudaProvider(); if (providers.empty()) GTEST_SKIP(); OpTester test("SigmoidGateMul", 1, kMSDomain); From b679f3b0305ae997b832cbf85f2deb551c63c85f Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 22 Oct 2026 05:43:25 +0000 Subject: [PATCH 05/19] Add shared-expert preparation and projected sigmoid gates Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/contrib_ops/cuda/sigmoid_gate_mul.md | 25 +++++ .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 40 +++++++- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 94 +++++++++++++++++++ .../contrib_ops/cuda/bert/fused_tensor_ops.h | 9 ++ .../core/graph/contrib_ops/bert_defs.cc | 3 + 5 files changed, 169 insertions(+), 2 deletions(-) create mode 100644 docs/contrib_ops/cuda/sigmoid_gate_mul.md diff --git a/docs/contrib_ops/cuda/sigmoid_gate_mul.md b/docs/contrib_ops/cuda/sigmoid_gate_mul.md new file mode 100644 index 0000000000000..7e57d370f77ec --- /dev/null +++ b/docs/contrib_ops/cuda/sigmoid_gate_mul.md @@ -0,0 +1,25 @@ +# Projected sigmoid gates and shared-expert preparation + +The CUDA BF16 `SigmoidGateMul` operator retains its existing packed, broadcast +and SiLU modes. Two appended optional inputs support shared-expert fusion: + +- Input 3, `prefill_numerics`, is a CPU INT32 tensor with one value, 0 or 1. + Mode 1 selects BF16 sigmoid intermediates; omitting it preserves the existing + `sigmoid_in_fp32` attribute. +- Input 4, `gate_weight`, is BF16 `[input_width, 1]`. It requires + `broadcast_gate=1` and `silu_gate=0`. In this mode input 1 is the projection + activation with the same leading dimensions as input 0, rather than an + already projected scalar gate. Its final dimension is `input_width`. + +The projected-gate implementation computes the dot product in FP32 using a +256-thread CUB reduction, rounds to BF16, applies the selected sigmoid +arithmetic, and preserves the existing BF16 multiply/residual rounding. +The residual, when supplied, must match input 0. + +`LaunchPrepareSharedExpert` shares the same projected-gate reduction and BF16 +sigmoid arithmetic. One launch prepares the SiLU/up activation and the scalar +row gates for `MatMulBlockQuantizedFp8Weight(shared_expert=1)` in companion +PR #33473. No behavior changes occur without the new inputs or opt-in mode. + +The shared-expert native tests compare these paths bit-exactly with the +unfused CUDA operator sequence, including prefill and fallback dispatch. diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc index f2cc3aac368f0..efca7bb0b6ad9 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -46,7 +46,9 @@ ONNX_OPERATOR_KERNEL_EX( ONNX_OPERATOR_KERNEL_EX( SigmoidGateMul, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) + .InputMemoryType(OrtMemTypeCPUInput, 3), SigmoidGateMul); ONNX_OPERATOR_KERNEL_EX( @@ -152,12 +154,46 @@ SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) } Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { + const Tensor* numerics = context->Input(3); + 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; const Tensor* attention = context->Input(0); const Tensor* projection = context->Input(1); const Tensor* residual = context->Input(2); + const Tensor* gate_weight = context->Input(4); const int64_t count = attention->Shape().Size(); ORT_RETURN_IF_NOT(residual == nullptr || residual->Shape() == attention->Shape(), "SigmoidGateMul residual must have the same shape as attention"); + if (gate_weight != nullptr) { + ORT_RETURN_IF_NOT(broadcast_gate_ && !silu_gate_, + "SigmoidGateMul gate_weight requires broadcast_gate without silu_gate"); + const auto& shape = attention->Shape(); + const auto& activation_shape = projection->Shape(); + ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1 && shape[shape.NumDimensions() - 1] > 0 && + activation_shape.NumDimensions() == shape.NumDimensions() && + activation_shape[activation_shape.NumDimensions() - 1] > 0, + "SigmoidGateMul projected gate requires matching ranks and positive final dimensions"); + for (size_t axis = 0; axis + 1 < shape.NumDimensions(); ++axis) { + ORT_RETURN_IF_NOT(shape[axis] == activation_shape[axis], + "SigmoidGateMul projected gate leading dimensions must match attention"); + } + const int64_t input_width = activation_shape[activation_shape.NumDimensions() - 1]; + ORT_RETURN_IF_NOT(gate_weight->Shape().NumDimensions() == 2 && + gate_weight->Shape()[0] == input_width && gate_weight->Shape()[1] == 1, + "SigmoidGateMul gate_weight must have shape [input_width, 1]"); + Tensor* output = context->Output(0, shape); + return LaunchProjectedSigmoidGateMul( + Stream(context), reinterpret_cast(attention->Data()), + reinterpret_cast(projection->Data()), + reinterpret_cast(gate_weight->Data()), + residual == nullptr ? nullptr : reinterpret_cast(residual->Data()), + reinterpret_cast(output->MutableData()), + count / shape[shape.NumDimensions() - 1], input_width, shape[shape.NumDimensions() - 1], + dense_prefill ? false : sigmoid_in_fp32_); + } int64_t gate_row_width = 0; if (silu_gate_) { ORT_RETURN_IF_NOT(projection->Shape() == attention->Shape(), @@ -195,7 +231,7 @@ Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { reinterpret_cast(output->MutableData()), count, head_dim_, heads_, q_stride_, tok_stride_, q_offset_, gate_row_width, residual == nullptr ? nullptr : reinterpret_cast(residual->Data()), - sigmoid_in_fp32_, silu_gate_); + dense_prefill ? false : sigmoid_in_fp32_, silu_gate_); } UntileHyperConnectionWeights::UntileHyperConnectionWeights(const OpKernelInfo& info) diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu index 8878f442185b3..48cc0e7fe260f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -143,6 +144,67 @@ __device__ __forceinline__ float StableSigmoid(float value) { return exp_value / (1.0f + exp_value); } +__device__ __forceinline__ float RoundedBf16Sigmoid(float value, bool sigmoid_in_fp32) { + if (sigmoid_in_fp32) { + return __bfloat162float(__float2bfloat16_rn(StableSigmoid(value))); + } + using namespace onnxruntime::cuda; + const BFloat16 a(value); + return static_cast(a > BFloat16(0.0f) + ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) + : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a))); +} + +__device__ float ProjectedRowGate(const __nv_bfloat16* activation, const __nv_bfloat16* weight, + int64_t row, int64_t input_width, bool sigmoid_in_fp32) { + using Reduce = cub::BlockReduce; + __shared__ typename Reduce::TempStorage reduction; + __shared__ float gate; + float dot = 0.0f; + for (int64_t k = threadIdx.x; k < input_width; k += kThreads) { + dot = fmaf(__bfloat162float(activation[row * input_width + k]), + __bfloat162float(weight[k]), dot); + } + const float total = Reduce(reduction).Sum(dot); + if (threadIdx.x == 0) { + const float value = __bfloat162float(__float2bfloat16_rn(total)); + gate = RoundedBf16Sigmoid(value, sigmoid_in_fp32); + } + __syncthreads(); + return gate; +} + +__global__ void PrepareSharedExpertKernel( + const __nv_bfloat16* up, const __nv_bfloat16* silu_gate, const __nv_bfloat16* activation, + const __nv_bfloat16* weight, __nv_bfloat16* prepared, float* row_gates, + int width, int input_width) { + const int row = blockIdx.x; + const float gate = ProjectedRowGate(activation, weight, row, input_width, false); + if (threadIdx.x == 0) row_gates[row] = gate; + for (int column = threadIdx.x; column < width; column += blockDim.x) { + const int64_t index = static_cast(row) * width + column; + const float value = __bfloat162float(silu_gate[index]); + const auto silu = __float2bfloat16_rn(value * RoundedBf16Sigmoid(value, false)); + prepared[index] = __float2bfloat16_rn(__bfloat162float(up[index]) * __bfloat162float(silu)); + } +} + +__global__ void ProjectedSigmoidGateMulKernel( + const __nv_bfloat16* attention, const __nv_bfloat16* activation, + const __nv_bfloat16* weight, const __nv_bfloat16* residual, __nv_bfloat16* output, + int64_t input_width, int64_t output_width, bool sigmoid_in_fp32) { + const int64_t row = blockIdx.x; + const float gate = ProjectedRowGate(activation, weight, row, input_width, sigmoid_in_fp32); + for (int64_t column = threadIdx.x; column < output_width; column += kThreads) { + const int64_t index = row * output_width + column; + __nv_bfloat16 value = __float2bfloat16_rn(__bfloat162float(attention[index]) * gate); + if (residual != nullptr) { + value = __float2bfloat16_rn(__bfloat162float(value) + __bfloat162float(residual[index])); + } + output[index] = value; + } +} + template __global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention, const __nv_bfloat16* projection, @@ -264,6 +326,38 @@ Status LaunchRowSelect(cudaStream_t stream, const int64_t* selection, return CUDA_CALL(cudaGetLastError()); } +Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, + const uint16_t* activation, const uint16_t* weight, + uint16_t* prepared, float* row_gates, + int rows, int width, int input_width) { + ORT_RETURN_IF_NOT(rows >= 0 && width > 0 && input_width > 0, + "Shared expert preparation requires nonnegative rows and positive widths"); + if (rows == 0) return Status::OK(); + PrepareSharedExpertKernel<<>>( + reinterpret_cast(up), reinterpret_cast(silu_gate), + reinterpret_cast(activation), reinterpret_cast(weight), + reinterpret_cast<__nv_bfloat16*>(prepared), row_gates, width, input_width); + return CUDA_CALL(cudaGetLastError()); +} + +Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, + const uint16_t* activation, const uint16_t* weight, + const uint16_t* residual, uint16_t* output, + int64_t rows, int64_t input_width, int64_t output_width, + bool sigmoid_in_fp32) { + ORT_RETURN_IF_NOT(rows >= 0 && rows <= std::numeric_limits::max() && + input_width > 0 && output_width > 0, + "SigmoidGateMul projected gate dimensions exceed CUDA limits"); + if (rows == 0) return Status::OK(); + ProjectedSigmoidGateMulKernel<<(rows), kThreads, 0, stream>>>( + reinterpret_cast(attention), + reinterpret_cast(activation), + reinterpret_cast(weight), + reinterpret_cast(residual), reinterpret_cast<__nv_bfloat16*>(output), + input_width, output_width, sigmoid_in_fp32); + return CUDA_CALL(cudaGetLastError()); +} + Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, const uint16_t* projection, uint16_t* output, int64_t count, int head_dim, int heads, int q_stride, diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h index 46bb3fb756677..fe953604728a0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h @@ -27,6 +27,15 @@ Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, int64_t token_stride, int64_t q_offset, int64_t gate_row_width, const uint16_t* residual, bool sigmoid_in_fp32, bool silu_gate); +Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, + const uint16_t* activation, const uint16_t* weight, + const uint16_t* residual, uint16_t* output, + int64_t rows, int64_t input_width, int64_t output_width, + bool sigmoid_in_fp32); +Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, + const uint16_t* activation, const uint16_t* weight, + uint16_t* prepared, float* row_gates, + int rows, int width, int input_width); Status LaunchUntileHyperConnectionWeights(cudaStream_t stream, const uint16_t* input, uint16_t* output, int kind); diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index cf0e49d4d15fb..929b7d04d682a 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -5746,11 +5746,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Input(0, "attention", "BF16 attention values.", "T") .Input(1, "projection", "BF16 packed projection containing gate values.", "T") .Input(2, "residual", "Optional BF16 residual with the attention shape.", "T", OpSchema::Optional) + .Input(3, "prefill_numerics", "Optional CPU int32 scalar 0/1.", "TI", OpSchema::Optional) + .Input(4, "gate_weight", "Optional BF16 projected-gate weight [input_width,1], requires broadcast_gate.", "T", OpSchema::Optional) .Output(0, "output", "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 " "before optionally adding residual and rounding again.", "T") .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and output use BF16.") + .TypeConstraint("TI", {"tensor(int32)"}, "CPU numerical control.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); if (hasInputShape(ctx, 0)) updateOutputShape(ctx, 0, getInputShape(ctx, 0)); From b442b7c44486d2232b338055df8bb45500cea74d Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 22 Oct 2026 17:38:17 +0000 Subject: [PATCH 06/19] Apply lintrunner formatting to fused tensor operators Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 12 ++++++------ onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu index 48cc0e7fe260f..5f0850ce76ce9 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -151,12 +151,12 @@ __device__ __forceinline__ float RoundedBf16Sigmoid(float value, bool sigmoid_in using namespace onnxruntime::cuda; const BFloat16 a(value); return static_cast(a > BFloat16(0.0f) - ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) - : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a))); + ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) + : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a))); } __device__ float ProjectedRowGate(const __nv_bfloat16* activation, const __nv_bfloat16* weight, - int64_t row, int64_t input_width, bool sigmoid_in_fp32) { + int64_t row, int64_t input_width, bool sigmoid_in_fp32) { using Reduce = cub::BlockReduce; __shared__ typename Reduce::TempStorage reduction; __shared__ float gate; @@ -327,9 +327,9 @@ Status LaunchRowSelect(cudaStream_t stream, const int64_t* selection, } Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, - const uint16_t* activation, const uint16_t* weight, - uint16_t* prepared, float* row_gates, - int rows, int width, int input_width) { + const uint16_t* activation, const uint16_t* weight, + uint16_t* prepared, float* row_gates, + int rows, int width, int input_width) { ORT_RETURN_IF_NOT(rows >= 0 && width > 0 && input_width > 0, "Shared expert preparation requires nonnegative rows and positive widths"); if (rows == 0) return Status::OK(); diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h index fe953604728a0..f916a6997266b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h @@ -33,9 +33,9 @@ Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attent int64_t rows, int64_t input_width, int64_t output_width, bool sigmoid_in_fp32); Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, - const uint16_t* activation, const uint16_t* weight, - uint16_t* prepared, float* row_gates, - int rows, int width, int input_width); + const uint16_t* activation, const uint16_t* weight, + uint16_t* prepared, float* row_gates, + int rows, int width, int input_width); Status LaunchUntileHyperConnectionWeights(cudaStream_t stream, const uint16_t* input, uint16_t* output, int kind); From dbd8b83c3e954b7d0814e9a73e4ea68f76699ffd Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 22 Oct 2026 17:38:31 +0000 Subject: [PATCH 07/19] Remove trailing whitespace in hyperconnection operator documentation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index d59ba2f78aef7..8b359d55f8cc6 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -3496,20 +3496,20 @@ This version of the operator has been available since version 1 of the 'com.micr One hyper-connection block with a low-rank gate network, fused into one operator. C = num_branches streams of width H are kept as one flattened (..., C * H) tensor. With the optional prologue the operator first applies the post-mix of the previous block (identity stream mix): - + S[..., c, h] = streams[..., c, h] + prev_post_gate[..., c] * prev_block_output[..., h] - + and otherwise S = streams. It then computes, with `rms` taken per branch over H: - + hn[..., c, h] = S[..., c, h] / sqrt(mean_h(S[..., c, :]^2) + epsilon) * norm_scale[c, h] d = ScaledSiLU(hn @ down_weight, silu_alpha) (..., R) g[..., c, h] = sigmoid(d @ up_weight) per stream and feature pre_output = reduction_scale * sum_c(hn[..., c, h] * g[..., c, h]) (..., H) post_gate = sigmoid((hn @ inject_weight) / gate_div) * gate_mul (..., C) - + pre_output is the input of the next transformer block and post_gate is prev_post_gate of the next call. streams_out is S (the streams after the prologue; a copy of streams when there is no prologue). - + Numerics contract. Matrix products and reductions accumulate in float32. The result of every step is rounded to T once, as the output of the equivalent standard operator would be: S = cast_T(streams + prev_post_gate * prev_block_output), hn = cast_T(S * inv_rms * norm_scale), @@ -9156,4 +9156,3 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
- From 07091ebc3847084aea4e41643e1c664b68a3c335 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 22 Oct 2026 23:27:01 +0000 Subject: [PATCH 08/19] Extend EmbeddingRMSNorm, RowSelect, and deterministic HyperConnectionLowRankMix post-gate Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 43 +++++++- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 97 +++++++++++++++++-- .../contrib_ops/cuda/bert/fused_tensor_ops.h | 11 ++- .../hyper_connection_low_rank_mix_impl.cu | 47 +++++++-- .../core/graph/contrib_ops/bert_defs.cc | 21 +++- 5 files changed, 192 insertions(+), 27 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc index efca7bb0b6ad9..5960d457161d2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -32,6 +32,8 @@ ONNX_OPERATOR_KERNEL_EX( EmbeddingRMSNorm, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("I32", DataTypeImpl::GetTensorType()) + .TypeConstraint("TT", BuildKernelDefConstraints()) .TypeConstraint("T", DataTypeImpl::GetTensorType()) .TypeConstraint("F", DataTypeImpl::GetTensorType()), EmbeddingRMSNorm); @@ -91,20 +93,49 @@ Status EmbeddingRMSNorm::ComputeInternal(OpKernelContext* context) const { const Tensor* ids = context->Input(0); const Tensor* table = context->Input(1); const Tensor* weight = context->Input(2); + const Tensor* committed = context->InputCount() > 3 ? context->Input(3) : nullptr; ORT_RETURN_IF_NOT(table->Shape().NumDimensions() == 2, "EmbeddingRMSNorm table must have shape [vocabulary, hidden]"); + const int64_t vocab_size = table->Shape()[0]; const int64_t hidden_size = table->Shape()[1]; ORT_RETURN_IF_NOT(weight->Shape().NumDimensions() == 1 && weight->Shape()[0] == hidden_size, "EmbeddingRMSNorm weight must match the embedding hidden size"); + const int64_t rows = ids->Shape().Size(); TensorShapeVector output_dims(ids->Shape().GetDims().begin(), ids->Shape().GetDims().end()); output_dims.push_back(hidden_size); Tensor* output = context->Output(0, TensorShape(output_dims)); + int64_t* slot_ids_ptr = nullptr; + int32_t* seqlens_k_ptr = nullptr; + int32_t* cu_seqlens_ptr = nullptr; + if (context->OutputCount() > 1) { + Tensor* slot_ids = context->Output(1, TensorShape({1, rows})); + if (slot_ids != nullptr) slot_ids_ptr = slot_ids->MutableData(); + } + if (context->OutputCount() > 2) { + Tensor* seqlens_k = context->Output(2, TensorShape({1})); + if (seqlens_k != nullptr) seqlens_k_ptr = seqlens_k->MutableData(); + } + if (context->OutputCount() > 3) { + Tensor* cu_seqlens = context->Output(3, TensorShape({2})); + if (cu_seqlens != nullptr) cu_seqlens_ptr = cu_seqlens->MutableData(); + } + const int32_t* committed_ptr = committed != nullptr ? committed->Data() : nullptr; + if (table->IsDataType()) { + const int8_t* table_i8 = table->Data(); + const float* row_scales = reinterpret_cast(table_i8 + vocab_size * hidden_size); + return LaunchEmbeddingRMSNormInt8(Stream(context), ids->Data(), table_i8, row_scales, + weight->Data(), epsilon_, + reinterpret_cast(output->MutableData()), + rows, vocab_size, hidden_size, + committed_ptr, slot_ids_ptr, seqlens_k_ptr, cu_seqlens_ptr); + } return LaunchEmbeddingRMSNorm(Stream(context), ids->Data(), reinterpret_cast(table->Data()), weight->Data(), epsilon_, reinterpret_cast(output->MutableData()), - ids->Shape().Size(), table->Shape()[0], hidden_size); + rows, vocab_size, hidden_size, + committed_ptr, slot_ids_ptr, seqlens_k_ptr, cu_seqlens_ptr); } Status RowSelect::ComputeInternal(OpKernelContext* context) const { @@ -113,12 +144,14 @@ Status RowSelect::ComputeInternal(OpKernelContext* context) const { std::array inputs{}; std::array outputs{}; std::array row_sizes{}; + int64_t num_rows = 2; for (int i = 0; i < 3; ++i) { const Tensor* input = context->Input(i + 1); const auto& shape = input->Shape(); - ORT_RETURN_IF_NOT(shape.NumDimensions() >= 2 && shape[0] == 1 && shape[1] == 2, - "RowSelect inputs must have shape [1, 2, ...]"); - row_sizes[i] = shape.Size() / 2; + ORT_RETURN_IF_NOT(shape.NumDimensions() >= 2 && shape[0] == 1 && shape[1] >= 1, + "RowSelect inputs must have shape [1, R, ...]"); + num_rows = shape[1]; + row_sizes[i] = shape.Size() / shape[1]; TensorShapeVector output_dims(shape.GetDims().begin(), shape.GetDims().end()); output_dims[1] = 1; Tensor* output = context->Output(i, TensorShape(output_dims)); @@ -126,7 +159,7 @@ Status RowSelect::ComputeInternal(OpKernelContext* context) const { outputs[i] = reinterpret_cast(output->MutableData()); } return LaunchRowSelect(Stream(context), selection->Data(), inputs.data(), - outputs.data(), row_sizes.data()); + outputs.data(), row_sizes.data(), num_rows); } SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu index 5f0850ce76ce9..f5fa0f8592a19 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -85,8 +85,21 @@ Status LaunchTypedArgmax(cudaStream_t stream, const T* input, int64_t rows, __global__ void EmbedNormKernel(const int64_t* ids, const uint16_t* table, const float* weight, float epsilon, - uint16_t* output, int64_t vocab_size, int64_t hidden_size) { + uint16_t* output, int64_t vocab_size, int64_t hidden_size, + const int32_t* committed, int64_t* slot_ids, + int32_t* seqlens_k, int32_t* cu_seqlens, int64_t rows) { __shared__ float sums[512]; + if (threadIdx.x == 0 && committed != nullptr) { + const int32_t c = committed[0]; + if (slot_ids != nullptr) slot_ids[blockIdx.x] = static_cast(c) + static_cast(blockIdx.x); + if (blockIdx.x == 0) { + if (seqlens_k != nullptr) seqlens_k[0] = c + static_cast(rows) - 1; + if (cu_seqlens != nullptr) { + cu_seqlens[0] = 0; + cu_seqlens[1] = static_cast(rows); + } + } + } const int64_t row_index = blockIdx.x; int64_t token_id = ids[row_index]; if (token_id < -vocab_size || token_id >= vocab_size) { @@ -119,17 +132,68 @@ __global__ void EmbedNormKernel(const int64_t* ids, const uint16_t* table, } } +__global__ void EmbedNormKernelInt8(const int64_t* ids, const int8_t* table_i8, const float* row_scales, + const float* weight, float epsilon, + uint16_t* output, int64_t vocab_size, int64_t hidden_size, + const int32_t* committed, int64_t* slot_ids, + int32_t* seqlens_k, int32_t* cu_seqlens, int64_t rows) { + __shared__ float sums[512]; + if (threadIdx.x == 0 && committed != nullptr) { + const int32_t c = committed[0]; + if (slot_ids != nullptr) slot_ids[blockIdx.x] = static_cast(c) + static_cast(blockIdx.x); + if (blockIdx.x == 0) { + if (seqlens_k != nullptr) seqlens_k[0] = c + static_cast(rows) - 1; + if (cu_seqlens != nullptr) { + cu_seqlens[0] = 0; + cu_seqlens[1] = static_cast(rows); + } + } + } + const int64_t row_index = blockIdx.x; + int64_t token_id = ids[row_index]; + if (token_id < -vocab_size || token_id >= vocab_size) { + for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { + output[row_index * hidden_size + col] = __bfloat16_as_ushort(__float2bfloat16_rn(CUDART_NAN_F)); + } + return; + } + if (token_id < 0) token_id += vocab_size; + const int8_t* embedding = table_i8 + token_id * hidden_size; + const float rscale = row_scales[token_id]; + float sum = 0.0f; + for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { + const float value = __bfloat162float(__float2bfloat16_rn(static_cast(embedding[col]) * rscale)); + sum += value * value; + } + sums[threadIdx.x] = sum; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + sums[threadIdx.x] += sums[threadIdx.x + stride]; + } + __syncthreads(); + } + const float scale = 1.0f / sqrtf(sums[0] / static_cast(hidden_size) + epsilon); + for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { + const float value = __bfloat162float(__float2bfloat16_rn(static_cast(embedding[col]) * rscale)); + const float normalized = (value * scale) * weight[col]; + output[row_index * hidden_size + col] = + __bfloat16_as_ushort(__float2bfloat16_rn(normalized)); + } +} + __global__ void RowSelectKernel(const int64_t* selection, const uint16_t* input0, const uint16_t* input1, const uint16_t* input2, uint16_t* output0, uint16_t* output1, uint16_t* output2, int64_t row_size0, int64_t row_size1, - int64_t row_size2) { + int64_t row_size2, int64_t num_rows) { const int which = blockIdx.y; const uint16_t* input = which == 0 ? input0 : (which == 1 ? input1 : input2); uint16_t* output = which == 0 ? output0 : (which == 1 ? output1 : output2); const int64_t row_size = which == 0 ? row_size0 : (which == 1 ? row_size1 : row_size2); - const int64_t selected_row = selection[0] != 0 ? 1 : 0; + const int64_t raw_sel = selection[0]; + const int64_t selected_row = raw_sel <= 0 ? 0 : (raw_sel >= num_rows ? (num_rows - 1) : raw_sel); for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < row_size; i += static_cast(gridDim.x) * blockDim.x) { output[i] = input[selected_row * row_size + i]; @@ -305,24 +369,43 @@ Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_ty Status LaunchEmbeddingRMSNorm(cudaStream_t stream, const int64_t* ids, const uint16_t* table, const float* weight, float epsilon, uint16_t* output, - int64_t rows, int64_t vocab_size, int64_t hidden_size) { + int64_t rows, int64_t vocab_size, int64_t hidden_size, + const int32_t* committed, int64_t* slot_ids, + int32_t* seqlens_k, int32_t* cu_seqlens) { ORT_RETURN_IF_NOT(hidden_size > 0 && vocab_size > 0, "EmbeddingRMSNorm requires non-empty table dimensions"); ORT_RETURN_IF_NOT(rows <= std::numeric_limits::max(), "EmbeddingRMSNorm has too many rows for a CUDA grid"); if (rows == 0) return Status::OK(); EmbedNormKernel<<(rows), 512, 0, stream>>>( - ids, table, weight, epsilon, output, vocab_size, hidden_size); + ids, table, weight, epsilon, output, vocab_size, hidden_size, + committed, slot_ids, seqlens_k, cu_seqlens, rows); + return CUDA_CALL(cudaGetLastError()); +} + +Status LaunchEmbeddingRMSNormInt8(cudaStream_t stream, const int64_t* ids, const int8_t* table_i8, + const float* row_scales, const float* weight, float epsilon, + uint16_t* output, int64_t rows, int64_t vocab_size, int64_t hidden_size, + const int32_t* committed, int64_t* slot_ids, + int32_t* seqlens_k, int32_t* cu_seqlens) { + ORT_RETURN_IF_NOT(hidden_size > 0 && vocab_size > 0, + "EmbeddingRMSNorm requires non-empty table dimensions"); + ORT_RETURN_IF_NOT(rows <= std::numeric_limits::max(), + "EmbeddingRMSNorm has too many rows for a CUDA grid"); + if (rows == 0) return Status::OK(); + EmbedNormKernelInt8<<(rows), 512, 0, stream>>>( + ids, table_i8, row_scales, weight, epsilon, output, vocab_size, hidden_size, + committed, slot_ids, seqlens_k, cu_seqlens, rows); return CUDA_CALL(cudaGetLastError()); } Status LaunchRowSelect(cudaStream_t stream, const int64_t* selection, const uint16_t* const* inputs, uint16_t* const* outputs, - const int64_t* row_sizes) { + const int64_t* row_sizes, int64_t num_rows) { const dim3 grid(8, 3); RowSelectKernel<<>>( selection, inputs[0], inputs[1], inputs[2], outputs[0], outputs[1], outputs[2], - row_sizes[0], row_sizes[1], row_sizes[2]); + row_sizes[0], row_sizes[1], row_sizes[2], num_rows); return CUDA_CALL(cudaGetLastError()); } diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h index f916a6997266b..88ebac79bb6ee 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h @@ -17,10 +17,17 @@ inline int64_t GetArgMaxLastAxisPartitions(int64_t width) { } Status LaunchEmbeddingRMSNorm(cudaStream_t stream, const int64_t* ids, const uint16_t* table, const float* weight, float epsilon, uint16_t* output, - int64_t rows, int64_t vocab_size, int64_t hidden_size); + int64_t rows, int64_t vocab_size, int64_t hidden_size, + const int32_t* committed = nullptr, int64_t* slot_ids = nullptr, + int32_t* seqlens_k = nullptr, int32_t* cu_seqlens = nullptr); +Status LaunchEmbeddingRMSNormInt8(cudaStream_t stream, const int64_t* ids, const int8_t* table_i8, + const float* row_scales, const float* weight, float epsilon, + uint16_t* output, int64_t rows, int64_t vocab_size, int64_t hidden_size, + const int32_t* committed = nullptr, int64_t* slot_ids = nullptr, + int32_t* seqlens_k = nullptr, int32_t* cu_seqlens = nullptr); Status LaunchRowSelect(cudaStream_t stream, const int64_t* selection, const uint16_t* const* inputs, uint16_t* const* outputs, - const int64_t* row_sizes); + const int64_t* row_sizes, int64_t num_rows = 2); Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, const uint16_t* projection, uint16_t* output, int64_t count, int head_dim, int heads, int q_stride, diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu index 7e6e2197f66fe..e0417524648c2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu @@ -656,13 +656,40 @@ __global__ void PreMixKernel(const T* __restrict__ hn, const T* __restrict__ up, } template -__global__ void PostGateKernel(const T* __restrict__ inject, T* __restrict__ gate, int64_t count, float gate_div, - float gate_mul) { - const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= count) return; - const float t = Rnd(to_float(inject[i]) / gate_div); - const float u = Rnd(SigmoidF32(t)); - gate[i] = from_float(u * gate_mul); +__global__ void FusedInjectPostGateKernel(const T* __restrict__ hn, const T* __restrict__ w_inject, + T* __restrict__ gate_out, int64_t rows, int branches, int64_t K, + float gate_div, float gate_mul) { + const int64_t row = blockIdx.x; + if (row >= rows) return; + const int tid = threadIdx.x; + const T* x_row = hn + row * K; + __shared__ float smem[8][8]; + float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + for (int64_t k = tid; k < K; k += kThreads) { + const float xv = to_float(x_row[k]); + const T* wk = w_inject + k * branches; +#pragma unroll + for (int b = 0; b < 8; ++b) { + if (b < branches) acc[b] = fmaf(xv, to_float(wk[b]), acc[b]); + } + } +#pragma unroll + for (int b = 0; b < 8; ++b) { + if (b < branches) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) acc[b] += __shfl_down_sync(0xffffffffu, acc[b], off); + if ((tid & 31) == 0) smem[tid >> 5][b] = acc[b]; + } + } + __syncthreads(); + if (tid < branches) { + float v = 0.0f; +#pragma unroll + for (int w = 0; w < 8; ++w) v += smem[w][tid]; + const float t = Rnd(Rnd(v) / gate_div); + const float u = Rnd(SigmoidF32(t)); + gate_out[row * branches + tid] = from_float(u * gate_mul); + } } template @@ -848,7 +875,7 @@ Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_ ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); ORT_RETURN_IF_ERROR(Gemm(cublas, rows, config.rank, K, hn, tensors.down_weight, down)); - ORT_RETURN_IF_ERROR(Gemm(cublas, rows, config.branches, K, hn, tensors.inject_weight, inject)); + (void)inject; const int64_t down_count = rows * config.rank; ScaledSiluKernel<<(BlocksFor(down_count)), kThreads, 0, stream>>>( @@ -862,8 +889,8 @@ Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_ PreMixKernel<<(BlocksFor(pre_count)), kThreads, 0, stream>>>( hn, up, tensors.pre_output, pre_count, config.branches, config.hidden, config.reduction_scale); ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); - PostGateKernel<<(BlocksFor(groups)), kThreads, 0, stream>>>( - inject, tensors.post_gate, groups, config.gate_div, config.gate_mul); + FusedInjectPostGateKernel<<(rows), kThreads, 0, stream>>>( + hn, tensors.inject_weight, tensors.post_gate, rows, config.branches, K, config.gate_div, config.gate_mul); return CUDA_CALL(cudaGetLastError()); } diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 92a07a302bc2b..639c67e3a8049 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -5771,14 +5771,29 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Embedding row indices with any shape. Negative indices in [-vocabulary, -1] " "count from the end. Out-of-range indices produce NaN output rows.", "I") - .Input(1, "table", "BF16 embedding matrix [vocabulary, hidden].", "T") + .Input(1, "table", "BF16 or INT8 embedding matrix [vocabulary, hidden].", "TT") .Input(2, "weight", "FP32 RMS normalization weight [hidden].", "F") + .Input(3, "committed_length", "Optional int32[1] committed length for head capture.", "I32", OpSchema::Optional) .Output(0, "normalized", "BF16 normalized embeddings [shape(ids), hidden].", "T") + .Output(1, "slot_ids", "Optional int64[1, R] slot ids.", "I", OpSchema::Optional) + .Output(2, "seqlens_k", "Optional int32[1] seqlens_k.", "I32", OpSchema::Optional) + .Output(3, "cumulative_sequence_lengths", "Optional int32[2] cumulative sequence lengths.", "I32", OpSchema::Optional) .TypeConstraint("I", {"tensor(int64)"}, "Indices use int64.") - .TypeConstraint("T", {"tensor(bfloat16)"}, "Embedding and result use BF16.") + .TypeConstraint("I32", {"tensor(int32)"}, "Committed length and sequence lengths use int32.") + .TypeConstraint("TT", {"tensor(bfloat16)", "tensor(int8)"}, "Embedding table uses BF16 or INT8.") + .TypeConstraint("T", {"tensor(bfloat16)"}, "Normalized result uses BF16.") .TypeConstraint("F", {"tensor(float)"}, "Normalization weight uses FP32.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 1, 0); + updateOutputElemType(ctx, 0, ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); + if (ctx.getNumOutputs() > 1 && ctx.hasOutput(1)) { + updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_INT64); + } + if (ctx.getNumOutputs() > 2 && ctx.hasOutput(2)) { + updateOutputElemType(ctx, 2, ONNX_NAMESPACE::TensorProto_DataType_INT32); + } + if (ctx.getNumOutputs() > 3 && ctx.hasOutput(3)) { + updateOutputElemType(ctx, 3, ONNX_NAMESPACE::TensorProto_DataType_INT32); + } if (!hasInputShape(ctx, 0) || !hasInputShape(ctx, 1)) return; const auto& ids_shape = getInputShape(ctx, 0); const auto& table_shape = getInputShape(ctx, 1); From 6c978adf2793e6fbe2fa0914e24de8adb15f3ded Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 22 Oct 2026 23:47:24 +0000 Subject: [PATCH 09/19] Remove HyperConnectionLowRankMix from PR #33470 (covered by PR #33384) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 89 -- docs/OperatorKernels.md | 1 - .../bert/hyper_connection_low_rank_mix.cc | 165 ---- .../cuda/bert/hyper_connection_low_rank_mix.h | 25 - .../hyper_connection_low_rank_mix_impl.cu | 913 ------------------ .../bert/hyper_connection_low_rank_mix_impl.h | 67 -- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 6 - .../core/graph/contrib_ops/bert_defs.cc | 107 -- onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 - .../hyper_connection_low_rank_mix_op_test.cc | 347 ------- .../test_hyper_connection_low_rank_mix.py | 311 ------ 11 files changed, 2033 deletions(-) delete mode 100644 onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc delete mode 100644 onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h delete mode 100644 onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu delete mode 100644 onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h delete mode 100644 onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc delete mode 100644 onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 8b359d55f8cc6..22e21aaf6832e 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -58,7 +58,6 @@ Do not modify directly.* * com.microsoft.GridSample * com.microsoft.GroupNorm * com.microsoft.GroupQueryAttention - * com.microsoft.HyperConnectionLowRankMix * com.microsoft.HyperConnectionPostMix * com.microsoft.HyperConnectionPreMix * com.microsoft.Inverse @@ -3491,94 +3490,6 @@ This version of the operator has been available since version 1 of the 'com.micr -### **com.microsoft.HyperConnectionLowRankMix** - - One hyper-connection block with a low-rank gate network, fused into one operator. C = num_branches streams of - width H are kept as one flattened (..., C * H) tensor. With the optional prologue the operator first applies the - post-mix of the previous block (identity stream mix): - - S[..., c, h] = streams[..., c, h] + prev_post_gate[..., c] * prev_block_output[..., h] - - and otherwise S = streams. It then computes, with `rms` taken per branch over H: - - hn[..., c, h] = S[..., c, h] / sqrt(mean_h(S[..., c, :]^2) + epsilon) * norm_scale[c, h] - d = ScaledSiLU(hn @ down_weight, silu_alpha) (..., R) - g[..., c, h] = sigmoid(d @ up_weight) per stream and feature - pre_output = reduction_scale * sum_c(hn[..., c, h] * g[..., c, h]) (..., H) - post_gate = sigmoid((hn @ inject_weight) / gate_div) * gate_mul (..., C) - - pre_output is the input of the next transformer block and post_gate is prev_post_gate of the next call. - streams_out is S (the streams after the prologue; a copy of streams when there is no prologue). - - Numerics contract. Matrix products and reductions accumulate in float32. The result of every step is rounded - to T once, as the output of the equivalent standard operator would be: - S = cast_T(streams + prev_post_gate * prev_block_output), hn = cast_T(S * inv_rms * norm_scale), - down = cast_T(hn @ down_weight), Z = cast_T(silu_alpha * down), SG = cast_T(sigmoid(Z)), hidden = cast_T(Z * SG), - up = cast_T(hidden @ up_weight), g = cast_T(sigmoid(up)), pre_output = cast_T(reduction_scale * sum_c(hn * g)), - inject = cast_T(hn @ inject_weight), post_gate = cast_T(cast_T(sigmoid(cast_T(inject / gate_div))) * gate_mul). - sigmoid is evaluated in float32 with expf and rounded once to T; it does not reproduce a Sigmoid kernel that - rounds to T after every elementary operation. The accumulation order of the matrix products is fixed, so the - results are deterministic, but they are not bitwise identical to cuBLAS. - -#### Version - -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 added to the mean square of each branch before the reciprocal square root.
-
gate_div : float (default is 1.0)
-
Divisor applied to the inject logits before the sigmoid. Must not be zero.
-
gate_mul : float (default is 1.0)
-
Multiplier applied to the sigmoid of the inject logits.
-
num_branches : int (required)
-
Number of residual streams C. The last dimension of streams is C * H.
-
reduction_scale : float (default is 1.0)
-
Multiplier applied to the sum over branches of pre_output.
-
silu_alpha : float (default is 1.0)
-
Scale of the ScaledSiLU between the down and up projections.
-
- -#### Inputs (5 - 7) - -
-
streams : T
-
Flattened streams (..., C * H). With the prologue these are the streams before the previous block's post-mix.
-
norm_scale : T
-
Per-branch RMS norm weight with shape (C * H).
-
down_weight : T
-
Down projection with shape (C * H, R), row-major [K, N].
-
up_weight : T
-
Up projection with shape (R, C * H), row-major [K, N].
-
inject_weight : T
-
Post-gate projection with shape (C * H, C), row-major [K, N].
-
prev_block_output (optional) : T
-
Optional output of the previous block with shape (..., H). Provide together with prev_post_gate.
-
prev_post_gate (optional) : T
-
Optional post-mix gate of the previous block with shape (..., C). Provide together with prev_block_output.
-
- -#### Outputs (2 - 3) - -
-
pre_output : T
-
Block input with shape (..., H).
-
post_gate : T
-
Post-mix gate for the next call with shape (..., C).
-
streams_out (optional) : T
-
Optional streams after the prologue, with the same shape as streams. Equal to streams when the prologue inputs are absent.
-
- -#### Type Constraints - -
-
T : tensor(float), tensor(float16), tensor(bfloat16)
-
Constrain all floating-point inputs and outputs to the same type.
-
- - ### **com.microsoft.HyperConnectionPostMix** Mixes existing streams and injects one branch output: diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 3b6fe65a4383d..cada9bc379ad7 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1133,7 +1133,6 @@ The **OpSet Version** column uses the following notation: |GridSample|*in* X:**T1**
*in* Grid:**T1**
*out* Y:**T2**|1+|**T1** = tensor(float)
**T2** = tensor(float)| |GroupNorm|*in* X:**T**
*in* gamma:**M**
*in* beta:**M**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |GroupQueryAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T_CACHE**
*in* past_value:**T_CACHE**
*in* seqlens_k:**M**
*in* total_sequence_length:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* position_ids:**tensor(int64)**
*in* attention_bias:**T**
*in* head_sink:**T**
*in* k_scale:**T_KV_SCALE**
*in* v_scale:**T_KV_SCALE**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*out* output:**T**
*out* present_key:**T_CACHE**
*out* present_value:**T_CACHE**
*out* output_qk:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)
**T_CACHE** = tensor(bfloat16), tensor(float16), tensor(float8e4m3fn), tensor(int8), tensor(uint8)
**T_KV_SCALE** = tensor(float)| -|HyperConnectionLowRankMix|*in* streams:**T**
*in* norm_scale:**T**
*in* down_weight:**T**
*in* up_weight:**T**
*in* inject_weight:**T**
*in* prev_block_output:**T**
*in* prev_post_gate:**T**
*out* pre_output:**T**
*out* post_gate:**T**
*out* streams_out:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |HyperConnectionPostMix|*in* streams:**T**
*in* block_output:**T**
*in* post_mix:**M**
*in* stream_mix:**M**
*out* output:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |HyperConnectionPreMix|*in* streams:**T**
*in* pre_mix:**M**
*out* output:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |Inverse|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc deleted file mode 100644 index 982698c00b222..0000000000000 --- a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc +++ /dev/null @@ -1,165 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h" - -#include - -#include "contrib_ops/cpu/hyper_connection_helper.h" -#include "contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h" -#include "core/providers/cuda/cuda_type_conversion.h" - -namespace onnxruntime::contrib::cuda { - -#define REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - HyperConnectionLowRankMix, kMSDomain, 1, T, kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ - HyperConnectionLowRankMix); - -REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(float) -REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(MLFloat16) -REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(BFloat16) - -#undef REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL - -template -HyperConnectionLowRankMix::HyperConnectionLowRankMix(const OpKernelInfo& info) - : CudaKernel(info), - num_branches_(info.GetAttrOrDefault("num_branches", 0)), - epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), - reduction_scale_(info.GetAttrOrDefault("reduction_scale", 1.0f)), - silu_alpha_(info.GetAttrOrDefault("silu_alpha", 1.0f)), - gate_div_(info.GetAttrOrDefault("gate_div", 1.0f)), - gate_mul_(info.GetAttrOrDefault("gate_mul", 1.0f)) { - ORT_ENFORCE(num_branches_ > 0, "HyperConnectionLowRankMix requires a positive num_branches"); - ORT_ENFORCE(gate_div_ != 0.0f, "gate_div must be non-zero"); -} - -template -Status HyperConnectionLowRankMix::ComputeInternal(OpKernelContext* context) const { - using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; - - const Tensor* streams = context->Input(0); - const Tensor* norm_scale = context->Input(1); - const Tensor* down_weight = context->Input(2); - const Tensor* up_weight = context->Input(3); - const Tensor* inject_weight = context->Input(4); - const Tensor* prev_block_output = context->Input(5); - const Tensor* prev_post_gate = context->Input(6); - - hyper_connection::StreamShape params; - ORT_RETURN_IF_ERROR(hyper_connection::ResolveStreamShape(streams->Shape(), num_branches_, params)); - ORT_RETURN_IF_NOT(params.branches <= std::numeric_limits::max() && - params.hidden <= std::numeric_limits::max() / params.branches, - "branch or hidden dimension is too large for CUDA"); - const int64_t branches = params.branches; - const int64_t hidden = params.hidden; - const int64_t features = branches * hidden; - - ORT_RETURN_IF_NOT(norm_scale->Shape().Size() == features && norm_scale->Shape().NumDimensions() == 1, - "norm_scale must have shape (num_branches * hidden)"); - ORT_RETURN_IF_NOT(down_weight->Shape().NumDimensions() == 2 && down_weight->Shape()[0] == features, - "down_weight must have shape (num_branches * hidden, rank)"); - const int64_t rank = down_weight->Shape()[1]; - ORT_RETURN_IF_NOT(rank > 0 && rank <= std::numeric_limits::max(), "rank must be positive"); - ORT_RETURN_IF_NOT(up_weight->Shape().NumDimensions() == 2 && up_weight->Shape()[0] == rank && - up_weight->Shape()[1] == features, - "up_weight must have shape (rank, num_branches * hidden)"); - ORT_RETURN_IF_NOT(inject_weight->Shape().NumDimensions() == 2 && inject_weight->Shape()[0] == features && - inject_weight->Shape()[1] == branches, - "inject_weight must have shape (num_branches * hidden, num_branches)"); - ORT_RETURN_IF_NOT((prev_block_output == nullptr) == (prev_post_gate == nullptr), - "prev_block_output and prev_post_gate must be provided together"); - - const bool has_prologue = prev_block_output != nullptr; - if (has_prologue) { - const auto& dims = streams->Shape().GetDims(); - TensorShapeVector expected_output(dims.begin(), dims.begin() + params.prefix_rank); - expected_output.push_back(hidden); - TensorShapeVector expected_gate(dims.begin(), dims.begin() + params.prefix_rank); - expected_gate.push_back(branches); - ORT_RETURN_IF_NOT(prev_block_output->Shape() == TensorShape(expected_output), - "prev_block_output must have shape (..., hidden)"); - ORT_RETURN_IF_NOT(prev_post_gate->Shape() == TensorShape(expected_gate), - "prev_post_gate must have shape (..., num_branches)"); - } - - Tensor* pre_output = context->Output(0, TensorShape(params.reduced_shape)); - TensorShapeVector gate_shape(params.reduced_shape.begin(), params.reduced_shape.end() - 1); - gate_shape.push_back(branches); - Tensor* post_gate = context->Output(1, TensorShape(gate_shape)); - Tensor* streams_out = context->Output(2, streams->Shape()); - - if (params.rows == 0) { - return Status::OK(); - } - - HyperConnectionLowRankMixConfig config; - config.rows = params.rows; - config.branches = static_cast(branches); - config.hidden = static_cast(hidden); - config.rank = static_cast(rank); - config.epsilon = epsilon_; - config.reduction_scale = reduction_scale_; - config.silu_alpha = silu_alpha_; - config.gate_div = gate_div_; - config.gate_mul = gate_mul_; - - IAllocatorUniquePtr streams_scratch; - HyperConnectionLowRankMixTensors tensors; - tensors.streams = reinterpret_cast(streams->Data()); - tensors.norm_scale = reinterpret_cast(norm_scale->Data()); - tensors.down_weight = reinterpret_cast(down_weight->Data()); - tensors.up_weight = reinterpret_cast(up_weight->Data()); - tensors.inject_weight = reinterpret_cast(inject_weight->Data()); - if (has_prologue) { - tensors.prev_block_output = reinterpret_cast(prev_block_output->Data()); - tensors.prev_post_gate = reinterpret_cast(prev_post_gate->Data()); - // The kernels need the post-mixed streams; use scratch memory if the graph does not consume them. - if (streams_out != nullptr) { - tensors.streams_out = reinterpret_cast(streams_out->MutableData()); - } else { - streams_scratch = GetScratchBuffer(static_cast(params.rows * features), GetComputeStream(context)); - tensors.streams_out = reinterpret_cast(streams_scratch.get()); - } - } else if (streams_out != nullptr) { - // Without a prologue the streams pass through unchanged. - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(streams_out->MutableDataRaw(), streams->DataRaw(), - streams->SizeInBytes(), cudaMemcpyDeviceToDevice, Stream(context))); - } - tensors.pre_output = reinterpret_cast(pre_output->MutableData()); - tensors.post_gate = reinterpret_cast(post_gate->MutableData()); - - // The fused kernels use 16-byte vector accesses on every tensor. - auto aligned16 = [](const void* ptr) { return ptr == nullptr || reinterpret_cast(ptr) % 16 == 0; }; - const bool aligned = aligned16(tensors.streams) && aligned16(tensors.norm_scale) && - aligned16(tensors.down_weight) && aligned16(tensors.up_weight) && - aligned16(tensors.inject_weight) && aligned16(tensors.prev_block_output) && - aligned16(tensors.streams_out); - - const auto& device_prop = GetDeviceProp(); - if (aligned && CanUseFusedHyperConnectionLowRankMix(device_prop, config)) { - auto scratch = GetScratchBuffer(HyperConnectionLowRankMixFusedScratchElements(config), - GetComputeStream(context)); - return LaunchHyperConnectionLowRankMixFused(device_prop, Stream(context), config, tensors, - scratch.get()); - } - - const size_t rows = static_cast(params.rows); - auto hn = GetScratchBuffer(rows * static_cast(features), GetComputeStream(context)); - auto down = GetScratchBuffer(rows * static_cast(rank), GetComputeStream(context)); - auto up = GetScratchBuffer(rows * static_cast(features), GetComputeStream(context)); - auto inject = GetScratchBuffer(rows * static_cast(branches), GetComputeStream(context)); - return LaunchHyperConnectionLowRankMixUnfused( - GetCublasHandle(context), Stream(context), config, tensors, reinterpret_cast(hn.get()), - reinterpret_cast(down.get()), reinterpret_cast(up.get()), - reinterpret_cast(inject.get())); -} - -template class HyperConnectionLowRankMix; -template class HyperConnectionLowRankMix; -template class HyperConnectionLowRankMix; - -} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h deleted file mode 100644 index 246f44afb4f8b..0000000000000 --- a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h +++ /dev/null @@ -1,25 +0,0 @@ -// 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 { - -template -class HyperConnectionLowRankMix final : public onnxruntime::cuda::CudaKernel { - public: - explicit HyperConnectionLowRankMix(const OpKernelInfo& info); - Status ComputeInternal(OpKernelContext* context) const override; - - private: - int64_t num_branches_; - float epsilon_; - float reduction_scale_; - float silu_alpha_; - float gate_div_; - float gate_mul_; -}; - -} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu deleted file mode 100644 index e0417524648c2..0000000000000 --- a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu +++ /dev/null @@ -1,913 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -// Fused low-rank hyper-connection block (see the HyperConnectionLowRankMix schema in bert_defs.cc). -// -// Fused path (rows <= kHyperConnectionLowRankMixMaxFusedRows, 16-bit T, 16-byte aligned shapes): two kernels. -// HcDownKernel grid (column tiles + 1, splits, row groups) -// [post-mix of the previous block] + branchwise RMS norm + down GEMV + inject GEMV. -// The K = C * H reduction is split in slices that each stay inside one branch; every CTA -// writes fp32 partial sums for its slice. No atomics, no clusters, no grid barrier. -// HcUpPreKernel grid (ceil(H / 16), 1, row groups) -// reduces the partials in a fixed order, then ScaledSiLU, up GEMV, sigmoid, pre-mix and the -// inject gate. -// Both kernels use ordinary stream-ordered launches: their weights may be outputs of preceding nodes. -// Every intermediate is rounded to T at the points where the unfused ONNX subgraph rounds; the GEMVs -// accumulate in fp32 in a fixed (deterministic) order. -// -// Unfused path (any shape, any T): norm kernel + cuBLAS GEMMs (fp32 accumulation) + elementwise kernels. - -#include "contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h" - -#include -#include -#include -#include -#include -#include - -#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" -#include "core/providers/cuda/cuda_common.h" - -namespace onnxruntime::contrib::cuda { -namespace { - -constexpr int kThreads = 256; -constexpr int kColTile = 16; // down columns per CTA of the down kernel (two 16-byte weight vectors) -constexpr int kUpTile = 16; // features per CTA of the up kernel -constexpr int kMaxSliceRows = 1280; // K rows handled by one CTA of the down kernel -constexpr int kWeightSlots = 10; // 16-byte weight vectors per thread held in shared memory -constexpr int kMaxVecIters = 2; // 16-byte vectors per thread in the norm pass -constexpr int kMaxFusedHidden = kMaxVecIters * kThreads * 8; -constexpr int kMaxFusedRank = 512; -constexpr int kMaxFusedBranches = 8; - -__device__ __forceinline__ float SigmoidF32(float v) { - if (v >= 0.0f) return 1.0f / (1.0f + expf(-v)); - const float e = expf(v); - return e / (1.0f + e); -} - -template -__device__ __forceinline__ float Rnd(float v) { - return to_float(from_float(v)); -} - -template -__device__ __forceinline__ float BitsToFloat(uint16_t bits) { - if constexpr (std::is_same_v) { - return __uint_as_float(static_cast(bits) << 16); - } else { - return __half2float(__ushort_as_half(bits)); - } -} - -template -__device__ __forceinline__ uint16_t FloatToBits(float v) { - if constexpr (std::is_same_v) { - return __bfloat16_as_ushort(__float2bfloat16(v)); - } else { - return __half_as_ushort(__float2half(v)); - } -} - -template -__device__ __forceinline__ void Unpack2(uint32_t u, float& a, float& b) { - a = BitsToFloat(static_cast(u & 0xffffu)); - b = BitsToFloat(static_cast(u >> 16)); -} - -template -__device__ __forceinline__ uint32_t Pack2(float a, float b) { - return static_cast(FloatToBits(a)) | (static_cast(FloatToBits(b)) << 16); -} - -template -__device__ __forceinline__ void Unpack8(const uint4& u, float (&f)[8]) { - Unpack2(u.x, f[0], f[1]); - Unpack2(u.y, f[2], f[3]); - Unpack2(u.z, f[4], f[5]); - Unpack2(u.w, f[6], f[7]); -} - -template -__device__ __forceinline__ uint4 Pack8(const float (&f)[8]) { - uint4 u; - u.x = Pack2(f[0], f[1]); - u.y = Pack2(f[2], f[3]); - u.z = Pack2(f[4], f[5]); - u.w = Pack2(f[6], f[7]); - return u; -} - -__device__ __forceinline__ void GridWait() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - asm volatile("griddepcontrol.wait;" ::: "memory"); -#endif -} - -__device__ __forceinline__ void GridTrigger() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); -#endif -} - -__device__ __forceinline__ void CpAsync16(void* smem, const void* gmem) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(static_cast(__cvta_generic_to_shared(smem))), - "l"(gmem) - : "memory"); -#else - *reinterpret_cast(smem) = *reinterpret_cast(gmem); -#endif -} - -__device__ __forceinline__ void CpCommit() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - asm volatile("cp.async.commit_group;" ::: "memory"); -#endif -} - -__device__ __forceinline__ void CpWaitAll() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - asm volatile("cp.async.wait_group 0;" ::: "memory"); -#endif -} - -struct DownParams { - int64_t rows; - int branches; - int hidden; - int rank; - int pieces; // slices per branch - int slice; // K rows per slice (multiple of 8) - int splits; // branches * pieces - int partial_stride; // rank + branches, rounded up to a multiple of 4 - float epsilon; -}; - -struct UpParams { - int64_t rows; - int branches; - int hidden; - int rank; - int splits; - int partial_stride; - float silu_alpha; - float reduction_scale; - float gate_div; - float gate_mul; -}; - -// CTA (j, s, z): K slice s (rows [k0, k0 + len) of the flattened stream, inside branch c) of the ROWS rows -// z * ROWS ... z * ROWS + ROWS - 1, which share every weight read. -// j < ceil(R / 16) owns down columns [16 j, 16 j + 16); the last j owns the C inject columns. -template -__global__ void __launch_bounds__(kThreads) - HcDownKernel(const T* __restrict__ streams, const T* __restrict__ prev_out, const T* __restrict__ prev_gate, - const T* __restrict__ norm_w, const T* __restrict__ w_down, const T* __restrict__ w_inject, - T* __restrict__ streams_out, float* __restrict__ partial, float* __restrict__ inv_rms_out, - DownParams p) { - const int j = blockIdx.x; - const int s = blockIdx.y; - const int64_t row0 = static_cast(blockIdx.z) * ROWS; - const int tid = threadIdx.x; - const int branches = p.branches; - const int hidden = p.hidden; - const int64_t K = static_cast(branches) * hidden; - const int c = s / p.pieces; - const int piece = s - c * p.pieces; - const int h0 = piece * p.slice; - const int len = min(p.slice, hidden - h0); - const int col_tiles = (p.rank + kColTile - 1) / kColTile; - const bool is_inject = (j == col_tiles); - const int64_t k0 = static_cast(c) * hidden + h0; - int64_t row_idx[ROWS]; - bool row_ok[ROWS]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - row_ok[r] = row0 + r < p.rows; - row_idx[r] = row_ok[r] ? row0 + r : p.rows - 1; // keep loads of unused rows in bounds - } - - __shared__ __align__(16) uint4 wsm[kWeightSlots * kThreads]; - __shared__ __align__(16) uint16_t xs[ROWS][kMaxSliceRows]; - __shared__ float red[8][ROWS][kColTile]; - __shared__ float sq_red[8][ROWS]; - - const int cg = tid & 1; - const int rp = tid >> 1; - const int col = j * kColTile + cg * 8; - // Let the up kernel start right away and prefetch its weights; its results are only read after GridWait(). - GridTrigger(); - - // Weights do not depend on any previous kernel: start streaming them into shared memory first. - if (!is_inject) { - if (col < p.rank) { - const T* base = w_down + k0 * p.rank + col; -#pragma unroll - for (int i = 0; i < kWeightSlots; ++i) { - const int row = rp + 128 * i; - if (row < len) CpAsync16(&wsm[i * kThreads + tid], base + static_cast(row) * p.rank); - } - } - } else { - const int chunks = len * branches / 8; - for (int q = tid; q < chunks; q += kThreads) CpAsync16(&wsm[q], w_inject + k0 * branches + q * 8); - } - CpCommit(); - - // Norm weights for the 8 consecutive features of the slice this thread normalizes below. - const bool norm_active = tid * 8 < len; - uint4 wv = make_uint4(0u, 0u, 0u, 0u); - if (norm_active) wv = *reinterpret_cast(norm_w + k0 + tid * 8); - - // Everything below consumes results of the previous kernel. - GridWait(); - - const int nvec = hidden / 8; - uint4 sv[ROWS][kMaxVecIters]; - uint4 bv[ROWS][kMaxVecIters]; - float gate_c[ROWS]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - const uint4* sp = reinterpret_cast(streams + row_idx[r] * K + static_cast(c) * hidden); -#pragma unroll - for (int i = 0; i < kMaxVecIters; ++i) { - const int qi = tid + kThreads * i; - sv[r][i] = qi < nvec ? sp[qi] : make_uint4(0u, 0u, 0u, 0u); - } - if constexpr (kPost) { - gate_c[r] = to_float(prev_gate[row_idx[r] * branches + c]); - const uint4* bp = reinterpret_cast(prev_out + row_idx[r] * hidden); -#pragma unroll - for (int i = 0; i < kMaxVecIters; ++i) { - const int qi = tid + kThreads * i; - bv[r][i] = qi < nvec ? bp[qi] : make_uint4(0u, 0u, 0u, 0u); - } - } - } - - float sum_sq[ROWS]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - sum_sq[r] = 0.0f; -#pragma unroll - for (int i = 0; i < kMaxVecIters; ++i) { - const int qi = tid + kThreads * i; - if (qi < nvec) { - const int h = qi * 8; - float v[8]; - Unpack8(sv[r][i], v); - if constexpr (kPost) { - float b[8]; - Unpack8(bv[r][i], b); -#pragma unroll - for (int m = 0; m < 8; ++m) v[m] = Rnd(fmaf(gate_c[r], b[m], v[m])); - } -#pragma unroll - for (int m = 0; m < 8; ++m) sum_sq[r] += v[m] * v[m]; - if (h >= h0 && h < h0 + len) { - const uint4 packed = Pack8(v); - *reinterpret_cast(&xs[r][h - h0]) = packed; - if constexpr (kPost) { - if (j == 0 && row_ok[r]) { - *reinterpret_cast(streams_out + row_idx[r] * K + static_cast(c) * hidden + h) = - packed; - } - } - } - } - } -#pragma unroll - for (int off = 16; off > 0; off >>= 1) sum_sq[r] += __shfl_xor_sync(0xffffffffu, sum_sq[r], off); - if ((tid & 31) == 0) sq_red[tid >> 5][r] = sum_sq[r]; - } - __syncthreads(); - float inv_rms[ROWS]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - float total = 0.0f; -#pragma unroll - for (int w = 0; w < 8; ++w) total += sq_red[w][r]; - inv_rms[r] = rsqrtf(total / static_cast(hidden) + p.epsilon); - if (j == 0 && piece == 0 && tid == 0 && row_ok[r]) inv_rms_out[row_idx[r] * branches + c] = inv_rms[r]; - if (norm_active) { - float x[8], w[8]; - Unpack8(*reinterpret_cast(&xs[r][tid * 8]), x); - Unpack8(wv, w); -#pragma unroll - for (int m = 0; m < 8; ++m) x[m] = x[m] * inv_rms[r] * w[m]; - *reinterpret_cast(&xs[r][tid * 8]) = Pack8(x); - } - } - CpWaitAll(); - __syncthreads(); - - if (!is_inject) { - float acc[ROWS][8]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { -#pragma unroll - for (int m = 0; m < 8; ++m) acc[r][m] = 0.0f; - } - if (col < p.rank) { -#pragma unroll - for (int i = 0; i < kWeightSlots; ++i) { - const int row = rp + 128 * i; - if (row < len) { - float wf[8]; - Unpack8(wsm[i * kThreads + tid], wf); -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - const float xv = BitsToFloat(xs[r][row]); -#pragma unroll - for (int m = 0; m < 8; ++m) acc[r][m] = fmaf(xv, wf[m], acc[r][m]); - } - } - } - } - // Lanes of a warp share cg = lane & 1 and differ in rp: reduce over lane bits 1..4. -#pragma unroll - for (int r = 0; r < ROWS; ++r) { -#pragma unroll - for (int m = 0; m < 8; ++m) { -#pragma unroll - for (int off = 2; off < 32; off <<= 1) acc[r][m] += __shfl_xor_sync(0xffffffffu, acc[r][m], off); - } - if ((tid & 31) < 2) { -#pragma unroll - for (int m = 0; m < 8; ++m) red[tid >> 5][r][cg * 8 + m] = acc[r][m]; - } - } - __syncthreads(); - if (tid < kColTile * ROWS) { - const int r = tid / kColTile; - const int t = tid % kColTile; - float v = 0.0f; -#pragma unroll - for (int q = 0; q < 8; ++q) v += red[q][r][t]; - const int out_col = j * kColTile + t; - if (out_col < p.rank && row_ok[r]) partial[(row_idx[r] * p.splits + s) * p.partial_stride + out_col] = v; - } - } else { - const T* wi = reinterpret_cast(wsm); - float a[ROWS][kMaxFusedBranches]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { -#pragma unroll - for (int m = 0; m < kMaxFusedBranches; ++m) a[r][m] = 0.0f; - } -#pragma unroll - for (int i = 0; i < 5; ++i) { - const int k = tid + kThreads * i; - if (k < len) { - float wf[kMaxFusedBranches]; -#pragma unroll - for (int m = 0; m < kMaxFusedBranches; ++m) wf[m] = m < branches ? to_float(wi[k * branches + m]) : 0.0f; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - const float xv = BitsToFloat(xs[r][k]); -#pragma unroll - for (int m = 0; m < kMaxFusedBranches; ++m) a[r][m] = fmaf(xv, wf[m], a[r][m]); - } - } - } -#pragma unroll - for (int r = 0; r < ROWS; ++r) { -#pragma unroll - for (int m = 0; m < kMaxFusedBranches; ++m) { -#pragma unroll - for (int off = 16; off > 0; off >>= 1) a[r][m] += __shfl_xor_sync(0xffffffffu, a[r][m], off); - if ((tid & 31) == 0) red[tid >> 5][r][m] = a[r][m]; - } - } - __syncthreads(); - if (tid < branches * ROWS) { - const int r = tid / branches; - const int t = tid % branches; - float v = 0.0f; -#pragma unroll - for (int q = 0; q < 8; ++q) v += red[q][r][t]; - if (row_ok[r]) partial[(row_idx[r] * p.splits + s) * p.partial_stride + p.rank + t] = v; - } - } -} - -// CTA (hb, 0, z) owns features [16 hb, 16 hb + 16) of every branch for ROWS rows. -template -__global__ void __launch_bounds__(kThreads) - HcUpPreKernel(const T* __restrict__ streams, const T* __restrict__ norm_w, const T* __restrict__ w_up, - const float* __restrict__ partial, const float* __restrict__ inv_rms_in, - T* __restrict__ pre_out, T* __restrict__ gate_out, UpParams p) { - const int hb = blockIdx.x; - const int64_t row0 = static_cast(blockIdx.z) * ROWS; - const int tid = threadIdx.x; - const int branches = BRANCHES == 0 ? p.branches : BRANCHES; - const int rank = RANK == 0 ? p.rank : RANK; - constexpr int kBranchStorage = BRANCHES == 0 ? kMaxFusedBranches : BRANCHES; - const int hidden = p.hidden; - const int64_t K = static_cast(branches) * hidden; - const int h0 = hb * kUpTile; - const int groups = 2 * branches; // (branch, 8-feature half) pairs; divides 32 - const int phases = kThreads / groups; // row phases - const int g = tid % groups; - const int c = g >> 1; - const int half_id = g & 1; - const int rp = tid / groups; - const int col = h0 + half_id * 8; - const bool valid = col < hidden; - const int chunk_rows = phases * kWeightSlots; - int64_t row_idx[ROWS]; - bool row_ok[ROWS]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - row_ok[r] = row0 + r < p.rows; - row_idx[r] = row_ok[r] ? row0 + r : p.rows - 1; - } - - __shared__ __align__(16) uint4 wsm[kWeightSlots * kThreads]; - __shared__ __align__(16) float sf[ROWS][RANK == 0 ? kMaxFusedRank : RANK]; - __shared__ float sg[ROWS][kBranchStorage * kUpTile]; - - // Prefetch the first chunk of up weights before waiting for the down kernel. - const T* wbase = w_up + static_cast(c) * hidden + col; - if (valid) { -#pragma unroll - for (int i = 0; i < kWeightSlots; ++i) { - const int row = rp + phases * i; - if (row < rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); - } - } - CpCommit(); - - // The next block's down kernel may start (and prefetch its weights) while this kernel waits. - GridTrigger(); - GridWait(); - - float xv[ROWS][kBranchStorage], inv[ROWS][kBranchStorage], nw[kBranchStorage]; - const int hl = tid; // feature owned by threads < 16 in the final pre-mix - const bool owner = tid < kUpTile && (h0 + hl) < hidden; - if (owner) { -#pragma unroll - for (int b = 0; b < kBranchStorage; ++b) { - if (b < branches) { - nw[b] = to_float(norm_w[static_cast(b) * hidden + h0 + hl]); -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - xv[r][b] = to_float(streams[row_idx[r] * K + static_cast(b) * hidden + h0 + hl]); - inv[r][b] = inv_rms_in[row_idx[r] * branches + b]; - } - } - } - } - - // Fixed-order reduction of the K-split partial sums, then ScaledSiLU; four columns per thread. - const int nq = rank / 4; - const int cq = tid * 4; - float4 va[ROWS]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) va[r] = make_float4(0.f, 0.f, 0.f, 0.f); - if (tid < nq) { - for (int sp0 = 0; sp0 < p.splits; sp0 += 8) { - float4 pq[8][ROWS]; -#pragma unroll - for (int u = 0; u < 8; ++u) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - pq[u][r] = sp0 + u < p.splits - ? __ldca(reinterpret_cast( - partial + (row_idx[r] * p.splits + sp0 + u) * p.partial_stride + cq)) - : make_float4(0.f, 0.f, 0.f, 0.f); - } - } -#pragma unroll - for (int u = 0; u < 8; ++u) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - va[r].x += pq[u][r].x; - va[r].y += pq[u][r].y; - va[r].z += pq[u][r].z; - va[r].w += pq[u][r].w; - } - } - } -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - const float v[4] = {va[r].x, va[r].y, va[r].z, va[r].w}; - float out[4]; -#pragma unroll - for (int m = 0; m < 4; ++m) { - const float z = Rnd(Rnd(v[m]) * p.silu_alpha); - out[m] = Rnd(z * Rnd(SigmoidF32(z))); - } - *reinterpret_cast(&sf[r][cq]) = make_float4(out[0], out[1], out[2], out[3]); - } - } - if (hb == 0 && tid < branches * ROWS) { - const int r = tid / branches; - const int t = tid % branches; - float v = 0.0f; - for (int sp = 0; sp < p.splits; ++sp) v += partial[(row_idx[r] * p.splits + sp) * p.partial_stride + rank + t]; - const float inject = Rnd(v); - const float u = Rnd(SigmoidF32(Rnd(inject / p.gate_div))); - if (row_ok[r]) gate_out[row_idx[r] * branches + t] = from_float(u * p.gate_mul); - } - __syncthreads(); - - float acc[ROWS][8]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { -#pragma unroll - for (int m = 0; m < 8; ++m) acc[r][m] = 0.0f; - } - for (int chunk0 = 0; chunk0 < rank; chunk0 += chunk_rows) { - if (chunk0 > 0 && valid) { -#pragma unroll - for (int i = 0; i < kWeightSlots; ++i) { - const int row = chunk0 + rp + phases * i; - if (row < rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); - } - CpCommit(); - } - CpWaitAll(); - if (valid) { -#pragma unroll - for (int i = 0; i < kWeightSlots; ++i) { - const int row = chunk0 + rp + phases * i; - if (row < rank) { - float wf[8]; - Unpack8(wsm[i * kThreads + tid], wf); -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - const float sv = sf[r][row]; -#pragma unroll - for (int m = 0; m < 8; ++m) acc[r][m] = fmaf(sv, wf[m], acc[r][m]); - } - } - } - } - } - // Lanes of a warp share g = lane % groups and differ in rp. - for (int off = groups; off < 32; off <<= 1) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { -#pragma unroll - for (int m = 0; m < 8; ++m) acc[r][m] += __shfl_xor_sync(0xffffffffu, acc[r][m], off); - } - } - __syncthreads(); // every thread is done with its weights: reuse that shared memory for the cross-warp sums - float* red = reinterpret_cast(wsm); // [8 warps][ROWS][groups * 8] - if ((tid & 31) < groups) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { -#pragma unroll - for (int m = 0; m < 8; ++m) red[((tid >> 5) * ROWS + r) * (groups * 8) + g * 8 + m] = acc[r][m]; - } - } - __syncthreads(); - if (tid < groups * 8) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - float v = 0.0f; -#pragma unroll - for (int q = 0; q < 8; ++q) v += red[(q * ROWS + r) * (groups * 8) + tid]; - const int gg = tid >> 3; - const int local = (gg & 1) * 8 + (tid & 7); - sg[r][(gg >> 1) * kUpTile + local] = Rnd(SigmoidF32(Rnd(v))); - } - } - __syncthreads(); - if (owner) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - float sum = 0.0f; -#pragma unroll - for (int b = 0; b < kBranchStorage; ++b) { - if (b < branches) { - const float n = Rnd(xv[r][b] * inv[r][b] * nw[b]); - sum = fmaf(n, sg[r][b * kUpTile + hl], sum); - } - } - if (row_ok[r]) pre_out[row_idx[r] * hidden + h0 + hl] = from_float(sum * p.reduction_scale); - } - } -} - -// ----------------------------------------------------------------------------- unfused path kernels - -template -__global__ void NormKernel(const T* __restrict__ streams, const T* __restrict__ prev_out, - const T* __restrict__ prev_gate, const T* __restrict__ norm_w, T* __restrict__ hn, - T* __restrict__ streams_out, int branches, int hidden, float epsilon) { - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage storage; - __shared__ float inv_rms; - const int64_t group = blockIdx.x; - const int64_t row = group / branches; - const int c = static_cast(group - row * branches); - const int64_t base = group * hidden; - float gate_c = 0.0f; - if constexpr (kPost) gate_c = to_float(prev_gate[group]); - - float sum_sq = 0.0f; - for (int h = threadIdx.x; h < hidden; h += kThreads) { - float v = to_float(streams[base + h]); - if constexpr (kPost) { - v = Rnd(fmaf(gate_c, to_float(prev_out[row * hidden + h]), v)); - streams_out[base + h] = from_float(v); - } - sum_sq += v * v; - } - const float total = BlockReduce(storage).Sum(sum_sq); - if (threadIdx.x == 0) inv_rms = rsqrtf(total / static_cast(hidden) + epsilon); - __syncthreads(); - for (int h = threadIdx.x; h < hidden; h += kThreads) { - float v = to_float(streams[base + h]); - if constexpr (kPost) v = Rnd(fmaf(gate_c, to_float(prev_out[row * hidden + h]), v)); - hn[base + h] = from_float(v * inv_rms * to_float(norm_w[static_cast(c) * hidden + h])); - } -} - -template -__global__ void ScaledSiluKernel(T* __restrict__ x, int64_t count, float alpha) { - const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= count) return; - const float z = Rnd(to_float(x[i]) * alpha); - const float sig = Rnd(SigmoidF32(z)); - x[i] = from_float(z * sig); -} - -template -__global__ void PreMixKernel(const T* __restrict__ hn, const T* __restrict__ up, T* __restrict__ pre, - int64_t count, int branches, int hidden, float reduction_scale) { - const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= count) return; - const int64_t row = i / hidden; - const int h = static_cast(i - row * hidden); - const int64_t K = static_cast(branches) * hidden; - float sum = 0.0f; - for (int b = 0; b < branches; ++b) { - const int64_t idx = row * K + static_cast(b) * hidden + h; - const float gate = Rnd(SigmoidF32(to_float(up[idx]))); - sum = fmaf(to_float(hn[idx]), gate, sum); - } - pre[i] = from_float(sum * reduction_scale); -} - -template -__global__ void FusedInjectPostGateKernel(const T* __restrict__ hn, const T* __restrict__ w_inject, - T* __restrict__ gate_out, int64_t rows, int branches, int64_t K, - float gate_div, float gate_mul) { - const int64_t row = blockIdx.x; - if (row >= rows) return; - const int tid = threadIdx.x; - const T* x_row = hn + row * K; - __shared__ float smem[8][8]; - float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; - for (int64_t k = tid; k < K; k += kThreads) { - const float xv = to_float(x_row[k]); - const T* wk = w_inject + k * branches; -#pragma unroll - for (int b = 0; b < 8; ++b) { - if (b < branches) acc[b] = fmaf(xv, to_float(wk[b]), acc[b]); - } - } -#pragma unroll - for (int b = 0; b < 8; ++b) { - if (b < branches) { -#pragma unroll - for (int off = 16; off > 0; off >>= 1) acc[b] += __shfl_down_sync(0xffffffffu, acc[b], off); - if ((tid & 31) == 0) smem[tid >> 5][b] = acc[b]; - } - } - __syncthreads(); - if (tid < branches) { - float v = 0.0f; -#pragma unroll - for (int w = 0; w < 8; ++w) v += smem[w][tid]; - const float t = Rnd(Rnd(v) / gate_div); - const float u = Rnd(SigmoidF32(t)); - gate_out[row * branches + tid] = from_float(u * gate_mul); - } -} - -template -constexpr cudaDataType_t CudaDataType() { - if constexpr (std::is_same_v) { - return CUDA_R_32F; - } else if constexpr (std::is_same_v) { - return CUDA_R_16F; - } else { - return CUDA_R_16BF; - } -} - -// Row-major C[m, n] = A[m, k] * B[k, n] with fp32 accumulation and T output. -template -Status Gemm(cublasHandle_t cublas, int64_t m, int64_t n, int64_t k, const T* a, const T* b, T* c) { - ORT_RETURN_IF_NOT(m <= std::numeric_limits::max() && n <= std::numeric_limits::max() && - k <= std::numeric_limits::max(), - "HyperConnectionLowRankMix GEMM dimension is too large"); - const float alpha = 1.0f; - const float beta = 0.0f; - constexpr cudaDataType_t type = CudaDataType(); - CUBLAS_RETURN_IF_ERROR(cublasGemmEx(cublas, CUBLAS_OP_N, CUBLAS_OP_N, static_cast(n), static_cast(m), - static_cast(k), &alpha, b, type, static_cast(n), a, type, - static_cast(k), &beta, c, type, static_cast(n), - CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); - return Status::OK(); -} - -int64_t BlocksFor(int64_t count) { return (count + kThreads - 1) / kThreads; } - -struct FusedGeometry { - int pieces; - int slice; - int splits; - int col_tiles; - int partial_stride; -}; - -FusedGeometry GetFusedGeometry(const HyperConnectionLowRankMixConfig& config) { - FusedGeometry g; - g.pieces = (config.hidden + kMaxSliceRows - 1) / kMaxSliceRows; - g.slice = ((config.hidden + g.pieces - 1) / g.pieces + 7) / 8 * 8; - g.pieces = (config.hidden + g.slice - 1) / g.slice; - g.splits = config.branches * g.pieces; - g.col_tiles = (config.rank + kColTile - 1) / kColTile; - g.partial_stride = (config.rank + config.branches + 3) / 4 * 4; // keeps every row of partials 16-byte aligned - return g; -} - -template -Status LaunchWithProgrammaticSerialization(bool use_pdl, Kernel kernel, dim3 grid, dim3 block, cudaStream_t stream, - Args... args) { -#if defined(CUDART_VERSION) && CUDART_VERSION >= 12030 - if (use_pdl) { - cudaLaunchAttribute attribute; - attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization; - attribute.val.programmaticStreamSerializationAllowed = 1; - cudaLaunchConfig_t launch_config = {}; - launch_config.gridDim = grid; - launch_config.blockDim = block; - launch_config.dynamicSmemBytes = 0; - launch_config.stream = stream; - launch_config.attrs = &attribute; - launch_config.numAttrs = 1; - return CUDA_CALL(cudaLaunchKernelEx(&launch_config, kernel, args...)); - } -#else - ORT_UNUSED_PARAMETER(use_pdl); -#endif - kernel<<>>(args...); - return CUDA_CALL(cudaGetLastError()); -} - -template -Status LaunchFusedRows(bool use_pdl, cudaStream_t stream, const HyperConnectionLowRankMixConfig& config, - const HyperConnectionLowRankMixTensors& tensors, const FusedGeometry& g, float* partial, - float* inv_rms, const DownParams& down, const UpParams& up) { - const unsigned row_groups = static_cast((config.rows + ROWS - 1) / ROWS); - const dim3 down_grid(g.col_tiles + 1, g.splits, row_groups); - const dim3 up_grid((config.hidden + kUpTile - 1) / kUpTile, 1, row_groups); - if (tensors.prev_block_output != nullptr) { - ORT_RETURN_IF_ERROR(LaunchWithProgrammaticSerialization( - use_pdl, HcDownKernel, down_grid, dim3(kThreads), stream, tensors.streams, - tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, tensors.down_weight, - tensors.inject_weight, tensors.streams_out, partial, inv_rms, down)); - } else { - ORT_RETURN_IF_ERROR(LaunchWithProgrammaticSerialization( - use_pdl, HcDownKernel, down_grid, dim3(kThreads), stream, tensors.streams, - static_cast(nullptr), static_cast(nullptr), tensors.norm_scale, - tensors.down_weight, tensors.inject_weight, static_cast(nullptr), partial, inv_rms, down)); - } - - const T* up_streams = tensors.prev_block_output != nullptr ? tensors.streams_out : tensors.streams; - if (config.branches == 4 && config.rank == 320) { - return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), stream, - up_streams, tensors.norm_scale, tensors.up_weight, - static_cast(partial), - static_cast(inv_rms), tensors.pre_output, - tensors.post_gate, up); - } - return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), stream, - up_streams, tensors.norm_scale, tensors.up_weight, - static_cast(partial), - static_cast(inv_rms), tensors.pre_output, - tensors.post_gate, up); -} - -} // namespace - -template -bool CanUseFusedHyperConnectionLowRankMix(const cudaDeviceProp& device_prop, - const HyperConnectionLowRankMixConfig& config) { - if constexpr (sizeof(T) != 2) { - return false; - } else { - const bool branches_ok = config.branches == 1 || config.branches == 2 || config.branches == 4 || - config.branches == 8; - return device_prop.major >= 8 && config.rows >= 1 && - config.rows <= kHyperConnectionLowRankMixMaxFusedRows && branches_ok && - config.hidden % 8 == 0 && config.hidden <= kMaxFusedHidden && - config.rank % 8 == 0 && config.rank >= 8 && config.rank <= kMaxFusedRank; - } -} - -size_t HyperConnectionLowRankMixFusedScratchElements(const HyperConnectionLowRankMixConfig& config) { - const FusedGeometry g = GetFusedGeometry(config); - return static_cast(config.rows) * - (static_cast(g.splits) * g.partial_stride + static_cast(config.branches)); -} - -template -Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, cudaStream_t stream, - const HyperConnectionLowRankMixConfig& config, - const HyperConnectionLowRankMixTensors& tensors, - float* scratch) { - if constexpr (sizeof(T) != 2) { - ORT_UNUSED_PARAMETER(device_prop); - ORT_UNUSED_PARAMETER(stream); - ORT_UNUSED_PARAMETER(config); - ORT_UNUSED_PARAMETER(tensors); - ORT_UNUSED_PARAMETER(scratch); - return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "fused HyperConnectionLowRankMix requires a 16-bit type"); - } else { - ORT_RETURN_IF_NOT(CanUseFusedHyperConnectionLowRankMix(device_prop, config), - "unsupported problem for the fused HyperConnectionLowRankMix path"); - const FusedGeometry g = GetFusedGeometry(config); - float* partial = scratch; - float* inv_rms = scratch + static_cast(config.rows) * g.splits * g.partial_stride; - - DownParams down{config.rows, config.branches, config.hidden, config.rank, g.pieces, g.slice, - g.splits, g.partial_stride, config.epsilon}; - UpParams up{config.rows, config.branches, config.hidden, config.rank, g.splits, g.partial_stride, - config.silu_alpha, config.reduction_scale, config.gate_div, config.gate_mul}; - // Weights can be produced by preceding nodes, so no input prefetch may overlap their execution. - const bool use_pdl = false; - if (config.rows == 1) { - return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); - } - return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); - } -} - -template -Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_t stream, - const HyperConnectionLowRankMixConfig& config, - const HyperConnectionLowRankMixTensors& tensors, - T* hn, T* down, T* up, T* inject) { - const int64_t rows = config.rows; - const int64_t K = static_cast(config.branches) * config.hidden; - const int64_t groups = rows * config.branches; - ORT_RETURN_IF_NOT(groups <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); - - if (tensors.prev_block_output != nullptr) { - NormKernel<<(groups), kThreads, 0, stream>>>( - tensors.streams, tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, hn, - tensors.streams_out, config.branches, config.hidden, config.epsilon); - } else { - NormKernel<<(groups), kThreads, 0, stream>>>( - tensors.streams, nullptr, nullptr, tensors.norm_scale, hn, nullptr, config.branches, config.hidden, - config.epsilon); - } - ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); - - ORT_RETURN_IF_ERROR(Gemm(cublas, rows, config.rank, K, hn, tensors.down_weight, down)); - (void)inject; - - const int64_t down_count = rows * config.rank; - ScaledSiluKernel<<(BlocksFor(down_count)), kThreads, 0, stream>>>( - down, down_count, config.silu_alpha); - ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); - ORT_RETURN_IF_ERROR(Gemm(cublas, rows, K, config.rank, down, tensors.up_weight, up)); - - const int64_t pre_count = rows * config.hidden; - ORT_RETURN_IF_NOT(BlocksFor(pre_count) <= std::numeric_limits::max(), - "CUDA launch requires too many blocks"); - PreMixKernel<<(BlocksFor(pre_count)), kThreads, 0, stream>>>( - hn, up, tensors.pre_output, pre_count, config.branches, config.hidden, config.reduction_scale); - ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); - FusedInjectPostGateKernel<<(rows), kThreads, 0, stream>>>( - hn, tensors.inject_weight, tensors.post_gate, rows, config.branches, K, config.gate_div, config.gate_mul); - return CUDA_CALL(cudaGetLastError()); -} - -#define INSTANTIATE(T) \ - template bool CanUseFusedHyperConnectionLowRankMix(const cudaDeviceProp&, \ - const HyperConnectionLowRankMixConfig&); \ - template Status LaunchHyperConnectionLowRankMixFused( \ - const cudaDeviceProp&, cudaStream_t, const HyperConnectionLowRankMixConfig&, \ - const HyperConnectionLowRankMixTensors&, float*); \ - template Status LaunchHyperConnectionLowRankMixUnfused( \ - cublasHandle_t, cudaStream_t, const HyperConnectionLowRankMixConfig&, \ - const HyperConnectionLowRankMixTensors&, T*, T*, T*, T*); - -INSTANTIATE(float) -INSTANTIATE(half) -INSTANTIATE(__nv_bfloat16) - -#undef INSTANTIATE - -} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h deleted file mode 100644 index c3c9c24c4842b..0000000000000 --- a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h +++ /dev/null @@ -1,67 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include -#include - -#include -#include - -#include "core/common/status.h" - -namespace onnxruntime::contrib::cuda { - -struct HyperConnectionLowRankMixConfig { - int64_t rows = 0; - int branches = 0; // C - int hidden = 0; // H, width of one branch; the streams have C * H features - int rank = 0; // R, width of the low-rank bottleneck - float epsilon = 1e-5f; - float reduction_scale = 1.0f; - float silu_alpha = 1.0f; - float gate_div = 1.0f; - float gate_mul = 1.0f; -}; - -template -struct HyperConnectionLowRankMixTensors { - const T* streams = nullptr; - const T* norm_scale = nullptr; - const T* down_weight = nullptr; - const T* up_weight = nullptr; - const T* inject_weight = nullptr; - const T* prev_block_output = nullptr; // prologue inputs; both null or both non-null - const T* prev_post_gate = nullptr; - T* pre_output = nullptr; - T* post_gate = nullptr; - T* streams_out = nullptr; // written when the prologue is present; must then be non-null -}; - -// Rows up to this value use the two-kernel fused path; larger inputs use the cuBLAS based path. -constexpr int64_t kHyperConnectionLowRankMixMaxFusedRows = 8; - -// True when the fused two-kernel path supports this problem on the given device. -template -bool CanUseFusedHyperConnectionLowRankMix(const cudaDeviceProp& device_prop, - const HyperConnectionLowRankMixConfig& config); - -// Number of float elements of scratch memory the fused path needs. -size_t HyperConnectionLowRankMixFusedScratchElements(const HyperConnectionLowRankMixConfig& config); - -template -Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, cudaStream_t stream, - const HyperConnectionLowRankMixConfig& config, - const HyperConnectionLowRankMixTensors& tensors, - float* scratch); - -// General path built from cuBLAS GEMMs and elementwise kernels. Scratch buffers (in elements of T): -// hn [rows, C * H], down [rows, R], up [rows, C * H], inject [rows, C]. -template -Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_t stream, - const HyperConnectionLowRankMixConfig& config, - const HyperConnectionLowRankMixTensors& tensors, - T* hn, T* down, T* up, T* inject); - -} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index de29a3579e58d..9e0c203e834c4 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -195,9 +195,6 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, HyperConnectionPreMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, HyperConnectionPostMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, HyperConnectionPostMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, HyperConnectionPostMix); -class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, HyperConnectionLowRankMix); -class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, HyperConnectionLowRankMix); -class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, HyperConnectionLowRankMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, EngramGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, EngramGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, EngramGate); @@ -569,9 +566,6 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 639c67e3a8049..cef02a698d5c9 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -6168,113 +6168,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( propagateShapeFromInputToOutput(ctx, 0, 0); })); -constexpr const char* HyperConnectionLowRankMix_ver1_doc = R"DOC( -One hyper-connection block with a low-rank gate network, fused into one operator. C = num_branches streams of -width H are kept as one flattened (..., C * H) tensor. With the optional prologue the operator first applies the -post-mix of the previous block (identity stream mix): - - S[..., c, h] = streams[..., c, h] + prev_post_gate[..., c] * prev_block_output[..., h] - -and otherwise S = streams. It then computes, with `rms` taken per branch over H: - - hn[..., c, h] = S[..., c, h] / sqrt(mean_h(S[..., c, :]^2) + epsilon) * norm_scale[c, h] - d = ScaledSiLU(hn @ down_weight, silu_alpha) (..., R) - g[..., c, h] = sigmoid(d @ up_weight) per stream and feature - pre_output = reduction_scale * sum_c(hn[..., c, h] * g[..., c, h]) (..., H) - post_gate = sigmoid((hn @ inject_weight) / gate_div) * gate_mul (..., C) - -pre_output is the input of the next transformer block and post_gate is prev_post_gate of the next call. -streams_out is S (the streams after the prologue; a copy of streams when there is no prologue). - -Numerics contract. Matrix products and reductions accumulate in float32. The result of every step is rounded -to T once, as the output of the equivalent standard operator would be: -S = cast_T(streams + prev_post_gate * prev_block_output), hn = cast_T(S * inv_rms * norm_scale), -down = cast_T(hn @ down_weight), Z = cast_T(silu_alpha * down), SG = cast_T(sigmoid(Z)), hidden = cast_T(Z * SG), -up = cast_T(hidden @ up_weight), g = cast_T(sigmoid(up)), pre_output = cast_T(reduction_scale * sum_c(hn * g)), -inject = cast_T(hn @ inject_weight), post_gate = cast_T(cast_T(sigmoid(cast_T(inject / gate_div))) * gate_mul). -sigmoid is evaluated in float32 with expf and rounded once to T; it does not reproduce a Sigmoid kernel that -rounds to T after every elementary operation. The accumulation order of the matrix products is fixed, so the -results are deterministic, but they are not bitwise identical to cuBLAS. -)DOC"; - -ONNX_MS_OPERATOR_SET_SCHEMA( - HyperConnectionLowRankMix, 1, - OpSchema() - .SetDoc(HyperConnectionLowRankMix_ver1_doc) - .Attr("num_branches", "Number of residual streams C. The last dimension of streams is C * H.", - AttributeProto::INT, true) - .Attr("epsilon", "Epsilon added to the mean square of each branch before the reciprocal square root.", - AttributeProto::FLOAT, 1e-5f) - .Attr("reduction_scale", "Multiplier applied to the sum over branches of pre_output.", - AttributeProto::FLOAT, 1.0f) - .Attr("silu_alpha", "Scale of the ScaledSiLU between the down and up projections.", - AttributeProto::FLOAT, 1.0f) - .Attr("gate_div", "Divisor applied to the inject logits before the sigmoid. Must not be zero.", - AttributeProto::FLOAT, 1.0f) - .Attr("gate_mul", "Multiplier applied to the sigmoid of the inject logits.", - AttributeProto::FLOAT, 1.0f) - .Input(0, "streams", - "Flattened streams (..., C * H). With the prologue these are the streams before the previous " - "block's post-mix.", - "T") - .Input(1, "norm_scale", "Per-branch RMS norm weight with shape (C * H).", "T") - .Input(2, "down_weight", "Down projection with shape (C * H, R), row-major [K, N].", "T") - .Input(3, "up_weight", "Up projection with shape (R, C * H), row-major [K, N].", "T") - .Input(4, "inject_weight", "Post-gate projection with shape (C * H, C), row-major [K, N].", "T") - .Input(5, "prev_block_output", - "Optional output of the previous block with shape (..., H). Provide together with prev_post_gate.", - "T", OpSchema::Optional) - .Input(6, "prev_post_gate", - "Optional post-mix gate of the previous block with shape (..., C). Provide together with " - "prev_block_output.", - "T", OpSchema::Optional) - .Output(0, "pre_output", "Block input with shape (..., H).", "T") - .Output(1, "post_gate", "Post-mix gate for the next call with shape (..., C).", "T") - .Output(2, "streams_out", - "Optional streams after the prologue, with the same shape as streams. Equal to streams when the " - "prologue inputs are absent.", - "T", OpSchema::Optional) - .TypeConstraint("T", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain all floating-point inputs and outputs to the same type.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - propagateElemTypeFromInputToOutput(ctx, 0, 1); - if (ctx.getNumOutputs() > 2) { - propagateElemTypeFromInputToOutput(ctx, 0, 2); - } - if (!hasInputShape(ctx, 0)) { - return; - } - const auto branches = getAttribute(ctx, "num_branches", 0); - if (branches < 1) { - fail_shape_inference("num_branches must be positive."); - } - const auto& input_shape = getInputShape(ctx, 0); - const int rank = input_shape.dim_size(); - if (rank < 1) { - fail_shape_inference("streams must have rank at least 1."); - } - auto* pre_shape = getOutputShape(ctx, 0); - auto* gate_shape = getOutputShape(ctx, 1); - for (int i = 0; i < rank - 1; ++i) { - *pre_shape->add_dim() = input_shape.dim(i); - *gate_shape->add_dim() = input_shape.dim(i); - } - auto* hidden = pre_shape->add_dim(); - const auto& last = input_shape.dim(rank - 1); - if (last.has_dim_value()) { - if (last.dim_value() % branches != 0) { - fail_shape_inference("The last dimension of streams must be divisible by num_branches."); - } - hidden->set_dim_value(last.dim_value() / branches); - } - gate_shape->add_dim()->set_dim_value(branches); - if (ctx.getNumOutputs() > 2) { - propagateShapeFromInputToOutput(ctx, 0, 2); - } - })); - constexpr const char* GatedAdd_ver1_doc = R"DOC( Adds one tensor to another tensor scaled by a per-row gate: diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index f43579623ac81..523f9d4a6a4c3 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -94,7 +94,6 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, 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); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionLowRankMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNet); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNetStateReplay); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ArgMaxLastAxis); @@ -235,7 +234,6 @@ 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/hyper_connection_low_rank_mix_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc deleted file mode 100644 index 65cf62aea7ebf..0000000000000 --- a/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc +++ /dev/null @@ -1,347 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include -#include -#include -#include -#include -#include - -#include "core/framework/session_options.h" -#include "core/graph/graph.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 -float RoundToType(float value) { - if constexpr (std::is_same_v) { - return MLFloat16(value).ToFloat(); - } else if constexpr (std::is_same_v) { - return BFloat16(value).ToFloat(); - } else { - return value; - } -} - -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; - } -} - -float Sigmoid(float v) { return 1.0f / (1.0f + std::exp(-v)); } - -struct Attrs { - float epsilon = 1e-6f; - float reduction_scale = 0.25f; - float silu_alpha = 0.25f; - float gate_div = 4.0f; - float gate_mul = 2.0f; -}; - -struct Problem { - int64_t rows, C, H, R; - std::vector streams, norm, down, up, inject, prev_out, prev_gate; -}; - -template -Problem MakeProblem(int64_t rows, int64_t C, int64_t H, int64_t R, bool prologue, unsigned seed) { - std::mt19937 gen(seed); - std::normal_distribution normal(0.0f, 1.0f); - const int64_t K = C * H; - auto fill = [&](size_t n, float scale, float shift) { - std::vector v(n); - for (auto& x : v) x = RoundToType(normal(gen) * scale + shift); - return v; - }; - Problem p{rows, C, H, R, {}, {}, {}, {}, {}, {}, {}}; - p.streams = fill(rows * K, 1.0f, 0.0f); - p.norm = fill(K, 0.1f, 1.0f); - p.down = fill(K * R, 1.0f / std::sqrt(static_cast(K)), 0.0f); - p.up = fill(R * K, 2.0f / std::sqrt(static_cast(R)), 0.0f); - p.inject = fill(K * C, 1.0f / std::sqrt(static_cast(K)), 0.0f); - if (prologue) { - p.prev_out = fill(rows * H, 1.0f, 0.0f); - p.prev_gate = fill(rows * C, 0.5f, 1.0f); - } - return p; -} - -struct Expected { - std::vector pre, gate, streams; -}; - -// float32 math with the rounding points listed in the schema documentation. -template -Expected Reference(const Problem& p, const Attrs& a) { - const int64_t K = p.C * p.H; - auto rd = [](float v) { return RoundToType(v); }; - Expected e; - e.streams = p.streams; - const bool prologue = !p.prev_out.empty(); - for (int64_t r = 0; r < p.rows; ++r) { - std::vector hn(K); - for (int64_t c = 0; c < p.C; ++c) { - float sum_sq = 0.0f; - for (int64_t h = 0; h < p.H; ++h) { - float v = p.streams[r * K + c * p.H + h]; - if (prologue) { - v = rd(v + p.prev_gate[r * p.C + c] * p.prev_out[r * p.H + h]); - e.streams[r * K + c * p.H + h] = v; - } - sum_sq += v * v; - } - const float inv = 1.0f / std::sqrt(sum_sq / static_cast(p.H) + a.epsilon); - for (int64_t h = 0; h < p.H; ++h) { - hn[c * p.H + h] = rd(e.streams[r * K + c * p.H + h] * inv * p.norm[c * p.H + h]); - } - } - std::vector hidden(p.R); - for (int64_t j = 0; j < p.R; ++j) { - double acc = 0.0; - for (int64_t k = 0; k < K; ++k) acc += static_cast(hn[k]) * p.down[k * p.R + j]; - const float z = rd(rd(static_cast(acc)) * a.silu_alpha); - hidden[j] = rd(z * rd(Sigmoid(z))); - } - std::vector gates(K); - for (int64_t k = 0; k < K; ++k) { - double acc = 0.0; - for (int64_t j = 0; j < p.R; ++j) acc += static_cast(hidden[j]) * p.up[j * K + k]; - gates[k] = rd(Sigmoid(rd(static_cast(acc)))); - } - for (int64_t h = 0; h < p.H; ++h) { - float sum = 0.0f; - for (int64_t c = 0; c < p.C; ++c) sum += hn[c * p.H + h] * gates[c * p.H + h]; - e.pre.push_back(rd(sum * a.reduction_scale)); - } - for (int64_t c = 0; c < p.C; ++c) { - double acc = 0.0; - for (int64_t k = 0; k < K; ++k) acc += static_cast(hn[k]) * p.inject[k * p.C + c]; - const float t = rd(rd(static_cast(acc)) / a.gate_div); - e.gate.push_back(rd(rd(Sigmoid(t)) * a.gate_mul)); - } - } - return e; -} - -template -void RunTest(int64_t rows, int64_t C, int64_t H, int64_t R, bool prologue, bool request_streams_out, - float tolerance, unsigned seed = 1) { - auto cuda = DefaultCudaExecutionProvider(); - if (!cuda) { - GTEST_SKIP() << "HyperConnectionLowRankMix is only implemented for CUDA"; - } - const Attrs a; - const Problem p = MakeProblem(rows, C, H, R, prologue, seed); - const Expected e = Reference(p, a); - const int64_t K = C * H; - - OpTester test("HyperConnectionLowRankMix", 1, kMSDomain); - test.AddAttribute("num_branches", C); - test.AddAttribute("epsilon", a.epsilon); - test.AddAttribute("reduction_scale", a.reduction_scale); - test.AddAttribute("silu_alpha", a.silu_alpha); - test.AddAttribute("gate_div", a.gate_div); - test.AddAttribute("gate_mul", a.gate_mul); - test.AddInput("streams", {rows, K}, ToTensorType(p.streams)); - test.AddInput("norm_scale", {K}, ToTensorType(p.norm)); - test.AddInput("down_weight", {K, R}, ToTensorType(p.down)); - test.AddInput("up_weight", {R, K}, ToTensorType(p.up)); - test.AddInput("inject_weight", {K, C}, ToTensorType(p.inject)); - if (prologue) { - test.AddInput("prev_block_output", {rows, H}, ToTensorType(p.prev_out)); - test.AddInput("prev_post_gate", {rows, C}, ToTensorType(p.prev_gate)); - } else { - test.AddOptionalInputEdge(); - test.AddOptionalInputEdge(); - } - test.AddOutput("pre_output", {rows, H}, ToTensorType(e.pre)); - test.AddOutput("post_gate", {rows, C}, ToTensorType(e.gate)); - if (request_streams_out) { - test.AddOutput("streams_out", {rows, K}, ToTensorType(e.streams)); - } else { - test.AddOptionalOutputEdge(); - } - test.SetOutputTolerance(tolerance, tolerance); - - std::vector> providers; - providers.push_back(std::move(cuda)); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - -constexpr float kFp32Tol = 2e-4f; -constexpr float kFp16Tol = 8e-3f; -constexpr float kBf16Tol = 6e-2f; - -class DependentNormScaleTester : public OpTester { - public: - DependentNormScaleTester() : OpTester("HyperConnectionLowRankMix", 1, kMSDomain) {} - - void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, - std::vector>&) override { - auto* norm = &graph.GetOrCreateNodeArg("produced_norm", nullptr); - auto* producer_gate = &graph.GetOrCreateNodeArg("producer_gate", nullptr); - const Attrs a; - auto add_mix = [&](const char* name, const std::vector& node_inputs, - const std::vector& node_outputs) { - auto& node = graph.AddNode(name, "HyperConnectionLowRankMix", "", node_inputs, node_outputs, nullptr, kMSDomain); - node.AddAttribute("num_branches", int64_t{1}); - node.AddAttribute("epsilon", a.epsilon); - node.AddAttribute("reduction_scale", a.reduction_scale); - node.AddAttribute("silu_alpha", a.silu_alpha); - node.AddAttribute("gate_div", a.gate_div); - node.AddAttribute("gate_mul", a.gate_mul); - }; - add_mix("producer", {inputs[0], inputs[1], inputs[2], inputs[3], inputs[4]}, {norm, producer_gate}); - add_mix("consumer", {inputs[5], norm, inputs[6], inputs[7], inputs[8]}, outputs); - graph.SetOutputs(outputs); - } -}; - -template -void RunDependentNormScaleTest(float tolerance) { - if (!DefaultCudaExecutionProvider()) GTEST_SKIP(); - for (unsigned seed : {17u, 29u, 53u}) { - const Attrs a; - const Problem producer = MakeProblem(1, 1, 64, 8, false, seed); - Problem consumer = MakeProblem(1, 1, 64, 8, false, seed + 1); - consumer.norm = Reference(producer, a).pre; - const Expected expected = Reference(consumer, a); - DependentNormScaleTester test; - test.AddInput("producer_streams", {64}, ToTensorType(producer.streams)); - test.AddInput("producer_norm", {64}, ToTensorType(producer.norm)); - test.AddInput("producer_down", {64, 8}, ToTensorType(producer.down)); - test.AddInput("producer_up", {8, 64}, ToTensorType(producer.up)); - test.AddInput("producer_inject", {64, 1}, ToTensorType(producer.inject)); - test.AddInput("consumer_streams", {64}, ToTensorType(consumer.streams)); - test.AddInput("consumer_down", {64, 8}, ToTensorType(consumer.down)); - test.AddInput("consumer_up", {8, 64}, ToTensorType(consumer.up)); - test.AddInput("consumer_inject", {64, 1}, ToTensorType(consumer.inject)); - test.AddOutput("pre_output", {64}, ToTensorType(expected.pre)); - test.AddOutput("post_gate", {1}, ToTensorType(expected.gate)); - test.SetOutputTolerance(tolerance, tolerance); - SessionOptions options; - options.graph_optimization_level = TransformerLevel::Default; - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - test.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); - } -} - -} // namespace - -TEST(HyperConnectionLowRankMixTest, FusedNormScaleProducedByPrecedingMix) { - RunDependentNormScaleTest(kFp16Tol); - RunDependentNormScaleTest(kBf16Tol); -} - -TEST(HyperConnectionLowRankMixTest, FusedFloat16WithPrologue) { - RunTest(2, 4, 16, 8, true, true, kFp16Tol); -} - -TEST(HyperConnectionLowRankMixTest, FusedFloat16WithoutPrologue) { - RunTest(1, 4, 64, 16, false, false, kFp16Tol); -} - -TEST(HyperConnectionLowRankMixTest, FusedBFloat16WithPrologueWithoutStreamsOut) { - RunTest(3, 4, 64, 16, true, false, kBf16Tol); -} - -TEST(HyperConnectionLowRankMixTest, FusedBFloat16TwoBranches) { - RunTest(2, 2, 32, 8, true, true, kBf16Tol); -} - -TEST(HyperConnectionLowRankMixTest, FusedMultipleKSlicesAndWeightChunks) { - RunTest(2, 4, 1288, 40, true, true, kBf16Tol); - RunTest(1, 4, 2560, 320, false, true, kFp16Tol); -} - -TEST(HyperConnectionLowRankMixTest, UnfusedManyRows) { - RunTest(9, 4, 16, 8, true, true, kFp16Tol); - RunTest(17, 4, 64, 16, false, true, kBf16Tol); -} - -TEST(HyperConnectionLowRankMixTest, FusedSpecializedRank320) { - for (int rows : {1, 2, 3}) { - RunTest(rows, 4, 1288, 320, true, true, kBf16Tol); - RunTest(rows, 4, 1288, 320, false, false, kFp16Tol); - } -} - -TEST(HyperConnectionLowRankMixTest, FusedDecodeShapeRank320) { - for (int rows : {1, 2}) { - for (bool prologue : {false, true}) { - RunTest(rows, 4, 2560, 320, prologue, true, kBf16Tol, 37); - RunTest(rows, 4, 2560, 320, prologue, true, kFp16Tol, 53); - } - } -} - -TEST(HyperConnectionLowRankMixTest, UnfusedUnsupportedShape) { - RunTest(2, 3, 10, 5, true, true, kFp16Tol); - RunTest(1, 4, 12, 16, false, false, kBf16Tol); -} - -TEST(HyperConnectionLowRankMixTest, Float32) { - RunTest(1, 4, 64, 16, true, true, kFp32Tol); - RunTest(5, 2, 24, 8, false, true, kFp32Tol); -} - -TEST(HyperConnectionLowRankMixTest, PrologueInputsMustBeProvidedTogether) { - auto cuda = DefaultCudaExecutionProvider(); - if (!cuda) { - GTEST_SKIP() << "HyperConnectionLowRankMix is only implemented for CUDA"; - } - constexpr int64_t C = 2, H = 8, R = 8, K = C * H; - OpTester test("HyperConnectionLowRankMix", 1, kMSDomain); - test.AddAttribute("num_branches", C); - test.AddInput("streams", {1, K}, std::vector(K, 1.0f)); - test.AddInput("norm_scale", {K}, std::vector(K, 1.0f)); - test.AddInput("down_weight", {K, R}, std::vector(K * R, 0.1f)); - test.AddInput("up_weight", {R, K}, std::vector(R * K, 0.1f)); - test.AddInput("inject_weight", {K, C}, std::vector(K * C, 0.1f)); - test.AddInput("prev_block_output", {1, H}, std::vector(H, 1.0f)); - test.AddOptionalInputEdge(); - test.AddOutput("pre_output", {1, H}, std::vector(H, 0.0f)); - test.AddOutput("post_gate", {1, C}, std::vector(C, 0.0f)); - test.AddOptionalOutputEdge(); - std::vector> providers; - providers.push_back(std::move(cuda)); - test.Run(OpTester::ExpectResult::kExpectFailure, "must be provided together", {}, nullptr, &providers); -} - -TEST(HyperConnectionLowRankMixTest, WeightShapeMismatchFails) { - auto cuda = DefaultCudaExecutionProvider(); - if (!cuda) { - GTEST_SKIP() << "HyperConnectionLowRankMix is only implemented for CUDA"; - } - constexpr int64_t C = 2, H = 8, R = 8, K = C * H; - OpTester test("HyperConnectionLowRankMix", 1, kMSDomain); - test.AddAttribute("num_branches", C); - test.AddInput("streams", {1, K}, std::vector(K, 1.0f)); - test.AddInput("norm_scale", {K}, std::vector(K, 1.0f)); - test.AddInput("down_weight", {K, R}, std::vector(K * R, 0.1f)); - test.AddInput("up_weight", {R + 1, K}, std::vector((R + 1) * K, 0.1f)); - test.AddInput("inject_weight", {K, C}, std::vector(K * C, 0.1f)); - test.AddOptionalInputEdge(); - test.AddOptionalInputEdge(); - test.AddOutput("pre_output", {1, H}, std::vector(H, 0.0f)); - test.AddOutput("post_gate", {1, C}, std::vector(C, 0.0f)); - test.AddOptionalOutputEdge(); - std::vector> providers; - providers.push_back(std::move(cuda)); - test.Run(OpTester::ExpectResult::kExpectFailure, "up_weight must have shape", {}, nullptr, &providers); -} - -} // namespace onnxruntime::test diff --git a/onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py b/onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py deleted file mode 100644 index 0709465b55125..0000000000000 --- a/onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py +++ /dev/null @@ -1,311 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -"""Numeric check for com.microsoft.HyperConnectionLowRankMix against a torch reference. - -The reference is the unfused subgraph the operator replaces, - - [HyperConnectionPostMix] -> BranchwiseRMSNorm -> MatMul(down) -> ScaledSiLU -> MatMul(up) -> Sigmoid - -> HyperConnectionPreMix, and MatMul(inject) -> Div -> Sigmoid -> Mul, - -evaluated with float32 math and rounded to the activation type at the points the standard operators round -(see the schema documentation), with a correctly rounded sigmoid. The fused path (few rows, 16-bit types) and -the cuBLAS based path (many rows, float32, shapes the fused kernels do not support) are both covered. -""" - -# Tensor dimensions keep the upper-case names the schema doc gives them. -# ruff: noqa: N806 - -import unittest - -import numpy as np -import torch -from onnx import TensorProto as TP # noqa: N817 -from onnx import helper - -import onnxruntime as ort - -DTYPES = { - TP.FLOAT: torch.float32, - TP.FLOAT16: torch.float16, - TP.BFLOAT16: torch.bfloat16, -} - -# (num_branches, hidden, rank): the real Qwen3.8-Flash-Next block first, then small and irregular shapes. -REAL_SHAPE = (4, 2560, 320) -SMALL_SHAPES = [ - (4, 64, 16), - (2, 32, 8), - (8, 64, 32), - (1, 48, 8), - (4, 1288, 40), # two K slices per branch in the fused path - (4, 4096, 512), # four K slices per branch and two weight chunks in the up kernel - (3, 40, 12), # not supported by the fused kernels (odd branch count, rank % 8 != 0) - (4, 100, 16), # hidden % 8 != 0 -] -ATTRS = dict(epsilon=1e-6, reduction_scale=0.25, silu_alpha=0.25, gate_div=4.0, gate_mul=2.0) - - -def has_cuda(): - return "CUDAExecutionProvider" in ort.get_available_providers() - - -def build_model(io_elem, C, H, R, prologue, attrs): - """Graph with float32 inputs/outputs; Casts move the activations to the type under test.""" - F = TP.FLOAT - K = C * H - node_inputs = ["streams", "norm_scale", "down_weight", "up_weight", "inject_weight"] - graph_inputs = [ - ("streams", ["rows", K]), - ("norm_scale", [K]), - ("down_weight", [K, R]), - ("up_weight", [R, K]), - ("inject_weight", [K, C]), - ] - if prologue: - node_inputs += ["prev_block_output", "prev_post_gate"] - graph_inputs += [("prev_block_output", ["rows", H]), ("prev_post_gate", ["rows", C])] - outputs = ["pre_output", "post_gate", "streams_out"] - - nodes = [helper.make_node("Cast", [n], [n + "_t"], to=io_elem, name=n + "/cast") for n, _ in graph_inputs] - nodes.append( - helper.make_node( - "HyperConnectionLowRankMix", - [n + "_t" for n in node_inputs], - [n + "_t" for n in outputs], - domain="com.microsoft", - num_branches=C, - **attrs, - ) - ) - nodes += [helper.make_node("Cast", [n + "_t"], [n], to=F, name=n + "/castback") for n in outputs] - graph = helper.make_graph( - nodes, - "hc_low_rank_mix", - [helper.make_tensor_value_info(n, F, shape) for n, shape in graph_inputs], - [ - helper.make_tensor_value_info("pre_output", F, ["rows", H]), - helper.make_tensor_value_info("post_gate", F, ["rows", C]), - helper.make_tensor_value_info("streams_out", F, ["rows", K]), - ], - ) - model = helper.make_model( - graph, opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)] - ) - model.ir_version = 10 - return model - - -def make_inputs(dtype, rows, C, H, R, prologue, seed): - g = torch.Generator().manual_seed(seed) - K = C * H - - def rnd(*shape, scale=1.0, shift=0.0): - return (torch.randn(*shape, generator=g) * scale + shift).to(dtype).float() - - inputs = { - "streams": rnd(rows, K), - "norm_scale": rnd(K, scale=0.1, shift=1.0), - "down_weight": rnd(K, R, scale=K**-0.5), - "up_weight": rnd(R, K, scale=2.0 * R**-0.5), - "inject_weight": rnd(K, C, scale=K**-0.5), - } - if prologue: - inputs["prev_block_output"] = rnd(rows, H) - inputs["prev_post_gate"] = (torch.rand(rows, C, generator=g) * 2.0).to(dtype).float() - return inputs - - -def reference(inputs, dtype, C, H, attrs): - def rd(t): - return t.to(dtype).float() - - def sigmoid(t): - return torch.sigmoid(t.double()).float() - - x = inputs["streams"] - rows = x.shape[0] - s = x.view(rows, C, H) - if "prev_block_output" in inputs: - s = rd(s + inputs["prev_post_gate"][:, :, None] * inputs["prev_block_output"][:, None, :]) - inv_rms = torch.rsqrt(s.pow(2).mean(-1, keepdim=True) + attrs["epsilon"]) - hn = rd(s * inv_rms * inputs["norm_scale"].view(C, H)) - hn_flat = hn.view(rows, C * H) - - down = rd(hn_flat @ inputs["down_weight"]) - z = rd(down * attrs["silu_alpha"]) - hidden = rd(z * rd(sigmoid(z))) - up = rd(hidden @ inputs["up_weight"]) - g = rd(sigmoid(up)).view(rows, C, H) - pre = rd(attrs["reduction_scale"] * (hn * g).sum(1)) - inject = rd(hn_flat @ inputs["inject_weight"]) - gate = rd(rd(sigmoid(rd(inject / attrs["gate_div"]))) * attrs["gate_mul"]) - return pre, gate, s.reshape(rows, C * H) - - -class SessionCache: - def __init__(self): - self.sessions = {} - - def get(self, io_elem, C, H, R, prologue, attrs): - key = (io_elem, C, H, R, prologue, tuple(sorted(attrs.items()))) - if key not in self.sessions: - model = build_model(io_elem, C, H, R, prologue, attrs) - so = ort.SessionOptions() - so.log_severity_level = 3 - self.sessions[key] = ort.InferenceSession( - model.SerializeToString(), so, providers=["CUDAExecutionProvider"] - ) - return self.sessions[key] - - -def tolerances(io_elem): - # One ulp of the output type plus slack for the different fp32 accumulation order flipping a rounding. - return {TP.FLOAT: (1e-3, 1e-5), TP.FLOAT16: (4e-3, 2e-3), TP.BFLOAT16: (3e-2, 8e-3)}[io_elem] - - -@unittest.skipUnless(has_cuda(), "HyperConnectionLowRankMix is a CUDA-only kernel") -class TestHyperConnectionLowRankMix(unittest.TestCase): - @classmethod - def setUpClass(cls): - cls.cache = SessionCache() - - def _check(self, io_elem, rows, shape, prologue, attrs=ATTRS, seed=0): - C, H, R = shape - dtype = DTYPES[io_elem] - inputs = make_inputs(dtype, rows, C, H, R, prologue, seed) - session = self.cache.get(io_elem, C, H, R, prologue, attrs) - pre, gate, streams_out = session.run(None, {k: v.numpy() for k, v in inputs.items()}) - exp_pre, exp_gate, exp_streams = reference(inputs, dtype, C, H, attrs) - rtol, atol = tolerances(io_elem) - msg = f"dtype={dtype} rows={rows} shape={shape} prologue={prologue}" - np.testing.assert_allclose(pre, exp_pre.numpy(), rtol=rtol, atol=atol, err_msg="pre_output " + msg) - np.testing.assert_allclose(gate, exp_gate.numpy(), rtol=rtol, atol=atol, err_msg="post_gate " + msg) - # The prologue is a single fp32 operation with one rounding: it matches exactly (up to fma contraction). - np.testing.assert_allclose(streams_out, exp_streams.numpy(), rtol=rtol, atol=atol, err_msg="streams_out " + msg) - return pre, gate, streams_out - - def test_real_shape(self): - for io_elem in DTYPES: - for rows in (1, 2, 4, 16, 128): - for prologue in (False, True): - with self.subTest(io_elem=io_elem, rows=rows, prologue=prologue): - self._check(io_elem, rows, REAL_SHAPE, prologue) - - def test_small_shapes(self): - for io_elem in DTYPES: - for shape in SMALL_SHAPES: - for rows in (1, 2, 16): - for prologue in (False, True): - with self.subTest(io_elem=io_elem, shape=shape, rows=rows, prologue=prologue): - self._check(io_elem, rows, shape, prologue, seed=1) - - def test_default_attributes(self): - attrs = dict(epsilon=1e-5, reduction_scale=1.0, silu_alpha=1.0, gate_div=1.0, gate_mul=1.0) - for io_elem in DTYPES: - for rows in (1, 3): - with self.subTest(io_elem=io_elem, rows=rows): - self._check(io_elem, rows, (4, 64, 16), True, attrs=attrs, seed=2) - - def test_other_gate_constants(self): - attrs = dict(epsilon=1e-6, reduction_scale=0.5, silu_alpha=0.75, gate_div=3.0, gate_mul=1.5) - for io_elem in (TP.FLOAT16, TP.BFLOAT16): - for rows in (1, 2): - with self.subTest(io_elem=io_elem, rows=rows): - self._check(io_elem, rows, REAL_SHAPE, True, attrs=attrs, seed=3) - - def test_chain_of_blocks(self): - """Two consecutive calls through the prologue equal the explicit post-mix between them.""" - C, H, R = REAL_SHAPE - for io_elem in (TP.FLOAT16, TP.BFLOAT16): - dtype = DTYPES[io_elem] - with self.subTest(io_elem=io_elem): - inputs = make_inputs(dtype, 2, C, H, R, False, seed=4) - first = self.cache.get(io_elem, C, H, R, False, ATTRS) - second = self.cache.get(io_elem, C, H, R, True, ATTRS) - _, gate, streams = first.run(None, {k: v.numpy() for k, v in inputs.items()}) - # Use the block "output" as an independent tensor (a stand-in for the attention/MLP result). - block_out = torch.randn(2, H, generator=torch.Generator().manual_seed(8)).to(dtype).float().numpy() - chained = dict(inputs) - chained["streams"] = streams - chained["prev_block_output"] = block_out - chained["prev_post_gate"] = gate - out = second.run(None, {k: np.asarray(v) for k, v in chained.items()}) - ref_inputs = {k: torch.from_numpy(np.asarray(v)) for k, v in chained.items()} - exp = reference(ref_inputs, dtype, C, H, ATTRS) - rtol, atol = tolerances(io_elem) - np.testing.assert_allclose(out[0], exp[0].numpy(), rtol=rtol, atol=atol) - np.testing.assert_allclose(out[1], exp[1].numpy(), rtol=rtol, atol=atol) - - def test_fused_rows_are_independent(self): - """The fused kernels use a fixed reduction order: a row gives the same bits at any batch size.""" - C, H, R = REAL_SHAPE - for io_elem in (TP.FLOAT16, TP.BFLOAT16): - dtype = DTYPES[io_elem] - with self.subTest(io_elem=io_elem): - inputs = make_inputs(dtype, 4, C, H, R, True, seed=5) - batch = self.cache.get(io_elem, C, H, R, True, ATTRS).run( - None, {k: v.numpy() for k, v in inputs.items()} - ) - for r in range(4): - row_inputs = { - k: (v[r : r + 1] if k in ("streams", "prev_block_output", "prev_post_gate") else v).numpy() - for k, v in inputs.items() - } - single = self.cache.get(io_elem, C, H, R, True, ATTRS).run(None, row_inputs) - for b, s in zip(batch, single, strict=True): - np.testing.assert_array_equal(b[r : r + 1], s) - - def test_no_prologue_streams_out_is_passthrough(self): - C, H, R = 4, 64, 16 - inputs = make_inputs(torch.bfloat16, 2, C, H, R, False, seed=6) - session = self.cache.get(TP.BFLOAT16, C, H, R, False, ATTRS) - _, _, streams_out = session.run(None, {k: v.numpy() for k, v in inputs.items()}) - np.testing.assert_array_equal(streams_out, inputs["streams"].numpy()) - - def test_cuda_graph_capture_and_replay(self): - """No host synchronization or unmanaged allocation: the op can be captured and replayed.""" - C, H, R = REAL_SHAPE - for io_elem in (TP.FLOAT16, TP.BFLOAT16): - for rows in (1, 2): - with self.subTest(io_elem=io_elem, rows=rows): - dtype = DTYPES[io_elem] - inputs = make_inputs(dtype, rows, C, H, R, True, seed=9) - model = build_model(io_elem, C, H, R, True, ATTRS) - so = ort.SessionOptions() - so.log_severity_level = 3 - graph_session = ort.InferenceSession( - model.SerializeToString(), - so, - providers=[("CUDAExecutionProvider", {"enable_cuda_graph": "1"})], - ) - binding = graph_session.io_binding() - keep = [] - for name, value in inputs.items(): - ortvalue = ort.OrtValue.ortvalue_from_numpy(value.numpy(), "cuda", 0) - keep.append(ortvalue) - binding.bind_ortvalue_input(name, ortvalue) - shapes = {"pre_output": (rows, H), "post_gate": (rows, C), "streams_out": (rows, C * H)} - outputs = {} - for name, shape in shapes.items(): - outputs[name] = ort.OrtValue.ortvalue_from_shape_and_type(shape, np.float32, "cuda", 0) - binding.bind_ortvalue_output(name, outputs[name]) - graph_session.run_with_iobinding(binding) # captures - graph_session.run_with_iobinding(binding) # replays - expected = self.cache.get(io_elem, C, H, R, True, ATTRS).run( - None, {k: v.numpy() for k, v in inputs.items()} - ) - for name, exp in zip(shapes, expected, strict=True): - np.testing.assert_array_equal(outputs[name].numpy(), exp, err_msg=name) - - def test_zero_rows(self): - C, H, R = 4, 64, 16 - inputs = make_inputs(torch.float16, 0, C, H, R, False, seed=7) - session = self.cache.get(TP.FLOAT16, C, H, R, False, ATTRS) - pre, gate, streams_out = session.run(None, {k: v.numpy() for k, v in inputs.items()}) - self.assertEqual(pre.shape, (0, H)) - self.assertEqual(gate.shape, (0, C)) - self.assertEqual(streams_out.shape, (0, C * H)) - - -if __name__ == "__main__": - unittest.main() From 3c16bb655ad91674b703ef6d91c9a22652c7f0f4 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Oct 2026 23:07:16 +0000 Subject: [PATCH 10/19] [CUDA] Remove unused HyperConnectionLowRankMix, UntileHyperConnectionWeights, EmbeddingRMSNorm, RowSelect, and EngramGate changes from PR #33470 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 150 +---------- docs/OperatorKernels.md | 5 +- .../contrib_ops/cpu/bert/engram_gate.cc | 25 +- .../contrib_ops/cuda/bert/engram_gate.cc | 27 +- .../contrib_ops/cuda/bert/engram_gate_impl.cu | 7 +- .../contrib_ops/cuda/bert/engram_gate_impl.h | 3 +- .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 117 --------- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 242 +++--------------- .../contrib_ops/cuda/bert/fused_tensor_ops.h | 45 +--- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 6 - .../core/graph/contrib_ops/bert_defs.cc | 122 +-------- onnxruntime/core/graph/contrib_ops/ms_opset.h | 6 - .../webgpu/webgpu_execution_provider.cc | 8 - .../test/contrib_ops/engram_ops_test.cc | 153 ++--------- .../test/contrib_ops/fused_tensor_ops_test.cc | 98 ------- 15 files changed, 108 insertions(+), 906 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 22e21aaf6832e..a08153206f497 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -35,7 +35,6 @@ Do not modify directly.* * com.microsoft.DynamicTimeWarping * com.microsoft.EPContext * com.microsoft.EmbedLayerNormalization - * com.microsoft.EmbeddingRMSNorm * com.microsoft.EngramGate * com.microsoft.ExpandDims * com.microsoft.FastGelu @@ -121,7 +120,6 @@ Do not modify directly.* * com.microsoft.RestorePadding * com.microsoft.Rfft * com.microsoft.RotaryEmbedding - * com.microsoft.RowSelect * com.microsoft.SampleOp * com.microsoft.Sampling * com.microsoft.ScaledSiLU @@ -140,7 +138,6 @@ Do not modify directly.* * com.microsoft.Trilu * com.microsoft.UnfoldTensor * com.microsoft.Unique - * com.microsoft.UntileHyperConnectionWeights * com.microsoft.VarlenCausalConvWithState * com.microsoft.VarlenNGramHashMapping * com.microsoft.WhisperBeamSearch @@ -2078,51 +2075,6 @@ This version of the operator has been available since version 1 of the 'com.micr -### **com.microsoft.EmbeddingRMSNorm** - - CUDA contrib operators for fused tensor operations used by transformer models. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
epsilon : float
-
RMS normalization epsilon.
-
- -#### Inputs - -
-
ids : I
-
Embedding row indices with any shape. Negative indices in [-vocabulary, -1] count from the end. Out-of-range indices produce NaN output rows.
-
table : T
-
BF16 embedding matrix [vocabulary, hidden].
-
weight : F
-
FP32 RMS normalization weight [hidden].
-
- -#### Outputs - -
-
normalized : T
-
BF16 normalized embeddings [shape(ids), hidden].
-
- -#### Type Constraints - -
-
I : tensor(int64)
-
Indices use int64.
-
T : tensor(bfloat16)
-
Embedding and result use BF16.
-
F : tensor(float)
-
Normalization weight uses FP32.
-
- - ### **com.microsoft.EngramGate** Fuses the Engram gate. @@ -2151,7 +2103,7 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
epsilon : float (default is 1e-05)
+
epsilon : float
Epsilon used by both RMS normalization steps. Default is 1e-5.
@@ -2159,11 +2111,11 @@ This version of the operator has been available since version 1 of the 'com.micr
key : T
-
Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size) or (total_tokens, hc_mult, hidden_size).
+
Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size).
query : T
-
Hidden-state queries with the same shape as key.
+
Hidden-state queries with shape (batch_size, sequence_length, hc_mult, hidden_size).
value : T
-
Projected Engram value shared by every hyper-connection, with shape (batch_size, sequence_length, hidden_size) or (total_tokens, hidden_size).
+
Projected Engram value shared by every hyper-connection, with shape (batch_size, sequence_length, hidden_size).
key_norm_scale : T
RMSNorm scale for keys with shape (hc_mult, hidden_size).
query_norm_scale : T
@@ -2176,9 +2128,9 @@ This version of the operator has been available since version 1 of the 'com.micr
output : T
-
Gated value tensor with the same shape as key.
+
Gated value tensor with shape (batch_size, sequence_length, hc_mult, hidden_size).
gated_value_normed (optional) : T
-
Optional RMS-normalized gated value tensor with shape (batch_size, sequence_length, hc_mult, hidden_size) or (total_tokens, hc_mult, hidden_size).
+
Optional RMS-normalized gated value tensor with shape (batch_size, sequence_length, hc_mult, hidden_size).
#### Type Constraints @@ -7383,48 +7335,6 @@ This version of the operator has been available since version 1 of the 'com.micr
-### **com.microsoft.RowSelect** - - CUDA contrib operators for fused tensor operations used by transformer models. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Inputs - -
-
selection : I
-
Scalar-like int64 row selector; zero selects row 0.
-
x0 : T
-
BF16 tensor with shape [1, 2, ...].
-
x1 : T
-
BF16 tensor with shape [1, 2, ...].
-
x2 : T
-
BF16 tensor with shape [1, 2, ...].
-
- -#### Outputs - -
-
y0 : T
-
Selected BF16 tensor with shape [1, 1, ...].
-
y1 : T
-
Selected BF16 tensor with shape [1, 1, ...].
-
y2 : T
-
Selected BF16 tensor with shape [1, 1, ...].
-
- -#### Type Constraints - -
-
I : tensor(int64)
-
Selector uses int64.
-
T : tensor(bfloat16)
-
Inputs and outputs use BF16.
-
- - ### **com.microsoft.SampleOp** Sample echo operator. @@ -7613,7 +7523,7 @@ This version of the operator has been available since version 1 of the 'com.micr
Per-token stride in the packed projection.
-#### Inputs (2 - 3) +#### Inputs (2 - 5)
attention : T
@@ -7622,6 +7532,10 @@ This version of the operator has been available since version 1 of the 'com.micr
BF16 packed projection containing gate values.
residual (optional) : T
Optional BF16 residual with the attention shape.
+
prefill_numerics (optional) : TI
+
Optional CPU int32 scalar 0/1.
+
gate_weight (optional) : T
+
Optional BF16 projected-gate weight [input_width,1], requires broadcast_gate.
#### Outputs @@ -7636,6 +7550,8 @@ This version of the operator has been available since version 1 of the 'com.micr
T : tensor(bfloat16)
Inputs and output use BF16.
+
TI : tensor(int32)
+
CPU numerical control.
@@ -8544,45 +8460,6 @@ This version of the operator has been available since version 1 of the 'com.micr -### **com.microsoft.UntileHyperConnectionWeights** - - CUDA contrib operators for fused tensor operations used by transformer models. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
kind : int
-
Untile layout kind, either 0 or 1.
-
- -#### Inputs (1 - 2) - -
-
tiled : T
-
Tiled BF16 weight [160, 20480].
-
unused (optional) : T
-
Reserved optional input retained for exporter compatibility.
-
- -#### Outputs - -
-
untiled : T
-
Untiled BF16 weight.
-
- -#### Type Constraints - -
-
T : tensor(bfloat16)
-
Input and output use BF16.
-
- - ### **com.microsoft.VarlenCausalConvWithState** Stateful causal depthwise convolution over a packed, token-major batch of variable-length @@ -9067,3 +8944,4 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
+ diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index cada9bc379ad7..ad597f81ae117 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1116,7 +1116,6 @@ The **OpSet Version** column uses the following notation: |DynamicSparseAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* auxiliary_key:**T**
*in* auxiliary_value:**T**
*in* selected_indices:**M**
*in* selected_counts:**M**
*in* seqlens_k:**M**
*in* total_sequence_length:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* position_ids:**tensor(int64)**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* head_sink:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |DynamicTimeWarping|*in* input:**F**
*out* output:**I**|1+|**F** = tensor(float)
**I** = tensor(int32)| |EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)| -|EmbeddingRMSNorm|*in* ids:**I**
*in* table:**T**
*in* weight:**F**
*out* normalized:**T**|1+|**F** = tensor(float)
**I** = tensor(int64)
**T** = tensor(bfloat16)| |EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)| @@ -1170,10 +1169,9 @@ The **OpSet Version** column uses the following notation: |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)| -|RowSelect|*in* selection:**I**
*in* x0:**T**
*in* x1:**T**
*in* x2:**T**
*out* y0:**T**
*out* y1:**T**
*out* y2:**T**|1+|**I** = tensor(int64)
**T** = tensor(bfloat16)| |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)| -|SigmoidGateMul|*in* attention:**T**
*in* projection:**T**
*in* residual:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16)| +|SigmoidGateMul|*in* attention:**T**
*in* projection:**T**
*in* residual:**T**
*in* prefill_numerics:**TI**
*in* gate_weight:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16)
**TI** = tensor(int32)| |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)| |SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |SkipSimplifiedLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| @@ -1183,7 +1181,6 @@ The **OpSet Version** column uses the following notation: |TransposeMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |Trilu|*in* X:**T**
*in* k:**tensor(int64)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |UnfoldTensor|*in* input:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| -|UntileHyperConnectionWeights|*in* tiled:**T**
*in* unused:**T**
*out* untiled:**T**|1+|**T** = tensor(bfloat16)| |VarlenCausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* cumulative_sequence_length:**M**
*in* bias:**T**
*in* initial_state:**T**
*in* capture_count:**M**
*out* output:**T**
*out* final_state:**T**
*out* state_update:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |VarlenNGramHashMapping|*in* input_ids:**M**
*in* multipliers:**M**
*in* vocab_sizes:**M**
*in* cumulative_sequence_length:**S**
*in* past_ids:**M**
*in* head_offsets:**M**
*in* eos_token_id:**M**
*in* segment_ids:**tensor(int32)**
*in* past_segment_ids:**S**
*in* capture_count:**S**
*out* hash_ids:**M**
*out* present_ids:**M**
*out* present_segment_ids:**S**
*out* state_update:**M**|1+|**M** = tensor(int32), tensor(int64)
**S** = tensor(int32)| |WhisperBeamSearch|*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**
*in* cross_qk_layer_head:**I**
*in* extra_decoding_ids:**I**
*in* temperature:**T**
*out* sequences:**I**
*out* sequences_scores:**T**
*out* scores:**T**
*out* cross_qk:**V**
*out* non_speech_probs:**T**|1+|**T** = tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc b/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc index edae2a9685d9f..c86c6005533cb 100644 --- a/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc +++ b/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc @@ -45,23 +45,16 @@ Status EngramGate::Compute(OpKernelContext* context) const { const Tensor* conv_norm_scale = context->Input(5); const TensorShape& key_shape = key->Shape(); - const size_t key_rank = key_shape.NumDimensions(); - ORT_RETURN_IF_NOT(key_rank == 3 || key_rank == 4, - "key must have shape (total_tokens, hc_mult, hidden_size) or " - "(batch_size, sequence_length, hc_mult, hidden_size)"); - const size_t token_dims = key_rank - 2; - const int64_t hc_mult = key_shape[token_dims]; - const int64_t hidden_size = key_shape[token_dims + 1]; - TensorShapeVector expected_value_dims; - expected_value_dims.reserve(key_rank - 1); - for (size_t i = 0; i < token_dims; ++i) { - expected_value_dims.push_back(key_shape[i]); - } - expected_value_dims.push_back(hidden_size); + ORT_RETURN_IF_NOT(key_shape.NumDimensions() == 4, + "key must have shape (batch_size, sequence_length, hc_mult, hidden_size)"); + const int64_t batch_size = key_shape[0]; + const int64_t sequence_length = key_shape[1]; + const int64_t hc_mult = key_shape[2]; + const int64_t hidden_size = key_shape[3]; ORT_RETURN_IF_NOT(query->Shape() == key_shape, "query must have the same shape as key"); - ORT_RETURN_IF_NOT(value->Shape() == TensorShape(expected_value_dims), - "value must match key's token dimensions and hidden_size"); + ORT_RETURN_IF_NOT(value->Shape() == TensorShape({batch_size, sequence_length, hidden_size}), + "value must have shape (batch_size, sequence_length, hidden_size)"); ORT_RETURN_IF_NOT(key_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}), "key_norm_scale must have shape (hc_mult, hidden_size)"); ORT_RETURN_IF_NOT(query_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}), @@ -88,7 +81,7 @@ Status EngramGate::Compute(OpKernelContext* context) const { T* output_data = output->MutableData(); T* output_normed_data = output_normed == nullptr ? nullptr : output_normed->MutableData(); - const int64_t rows = key_shape.SizeToDimension(key_rank - 1); + const int64_t rows = batch_size * sequence_length * hc_mult; ThreadPool::TryParallelFor( context->GetOperatorThreadPool(), narrow(rows), static_cast(2 * hidden_size + 32), diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc b/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc index e870950062761..825a35aba4121 100644 --- a/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc +++ b/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc @@ -45,24 +45,16 @@ Status EngramGate::ComputeInternal(OpKernelContext* context) const { const Tensor* conv_norm_scale = context->Input(5); const TensorShape& key_shape = key->Shape(); - const size_t key_rank = key_shape.NumDimensions(); - ORT_RETURN_IF_NOT(key_rank == 3 || key_rank == 4, - "key must have shape (total_tokens, hc_mult, hidden_size) or " - "(batch_size, sequence_length, hc_mult, hidden_size)"); - const size_t token_dims = key_rank - 2; - const int64_t token_count = key_shape.SizeToDimension(token_dims); - const int64_t hc_mult = key_shape[token_dims]; - const int64_t hidden_size = key_shape[token_dims + 1]; - TensorShapeVector expected_value_dims; - expected_value_dims.reserve(key_rank - 1); - for (size_t i = 0; i < token_dims; ++i) { - expected_value_dims.push_back(key_shape[i]); - } - expected_value_dims.push_back(hidden_size); + ORT_RETURN_IF_NOT(key_shape.NumDimensions() == 4, + "key must have shape (batch_size, sequence_length, hc_mult, hidden_size)"); + const int64_t batch_size = key_shape[0]; + const int64_t sequence_length = key_shape[1]; + const int64_t hc_mult = key_shape[2]; + const int64_t hidden_size = key_shape[3]; ORT_RETURN_IF_NOT(query->Shape() == key_shape, "query must have the same shape as key"); - ORT_RETURN_IF_NOT(value->Shape() == TensorShape(expected_value_dims), - "value must match key's token dimensions and hidden_size"); + ORT_RETURN_IF_NOT(value->Shape() == TensorShape({batch_size, sequence_length, hidden_size}), + "value must have shape (batch_size, sequence_length, hidden_size)"); ORT_RETURN_IF_NOT(key_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}), "key_norm_scale must have shape (hc_mult, hidden_size)"); ORT_RETURN_IF_NOT(query_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}), @@ -90,7 +82,8 @@ Status EngramGate::ComputeInternal(OpKernelContext* context) const { conv_norm_scale == nullptr ? nullptr : reinterpret_cast(conv_norm_scale->Data()), reinterpret_cast(output->MutableData()), output_normed == nullptr ? nullptr : reinterpret_cast(output_normed->MutableData()), - token_count, + batch_size, + sequence_length, hc_mult, hidden_size, epsilon_); diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu index fbcb761667430..788d9bf9ed72e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu @@ -96,11 +96,12 @@ Status LaunchEngramGateKernel( const T* conv_norm_scale, T* output, T* output_normed, - int64_t token_count, + int64_t batch_size, + int64_t sequence_length, int64_t hc_mult, int64_t hidden_size, float epsilon) { - const int64_t rows = token_count * hc_mult; + const int64_t rows = batch_size * sequence_length * hc_mult; if (rows == 0 || hidden_size == 0) { return Status::OK(); } @@ -114,7 +115,7 @@ Status LaunchEngramGateKernel( #define INSTANTIATE_ENGRAM_GATE(T) \ template Status LaunchEngramGateKernel(cudaStream_t, const T*, const T*, const T*, const T*, const T*, \ - const T*, T*, T*, int64_t, int64_t, int64_t, float); + const T*, T*, T*, int64_t, int64_t, int64_t, int64_t, float); INSTANTIATE_ENGRAM_GATE(float) INSTANTIATE_ENGRAM_GATE(half) diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h index 707a49daeb221..fdacbbd393600 100644 --- a/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h @@ -21,7 +21,8 @@ Status LaunchEngramGateKernel( const T* conv_norm_scale, T* output, T* output_normed, - int64_t token_count, + int64_t batch_size, + int64_t sequence_length, int64_t hc_mult, int64_t hidden_size, float epsilon); diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc index 5960d457161d2..1f360cf4d9b2a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -28,23 +28,6 @@ REGISTER_ARGMAX_LAST_AXIS(BFloat16) #undef REGISTER_ARGMAX_LAST_AXIS -ONNX_OPERATOR_KERNEL_EX( - EmbeddingRMSNorm, kMSDomain, 1, kCudaExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("I", DataTypeImpl::GetTensorType()) - .TypeConstraint("I32", DataTypeImpl::GetTensorType()) - .TypeConstraint("TT", BuildKernelDefConstraints()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()) - .TypeConstraint("F", DataTypeImpl::GetTensorType()), - EmbeddingRMSNorm); - -ONNX_OPERATOR_KERNEL_EX( - RowSelect, kMSDomain, 1, kCudaExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("I", DataTypeImpl::GetTensorType()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), - RowSelect); - ONNX_OPERATOR_KERNEL_EX( SigmoidGateMul, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) @@ -53,12 +36,6 @@ ONNX_OPERATOR_KERNEL_EX( .InputMemoryType(OrtMemTypeCPUInput, 3), SigmoidGateMul); -ONNX_OPERATOR_KERNEL_EX( - UntileHyperConnectionWeights, kMSDomain, 1, kCudaExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), - UntileHyperConnectionWeights); - template ArgMaxLastAxis::ArgMaxLastAxis(const OpKernelInfo& info) : CudaKernel(info), keepdims_(info.GetAttrOrDefault("keepdims", 0) != 0) {} @@ -84,84 +61,6 @@ Status ArgMaxLastAxis::ComputeInternal(OpKernelContext* context) const { output->MutableData(), partial_values.get(), partial_indices.get()); } -EmbeddingRMSNorm::EmbeddingRMSNorm(const OpKernelInfo& info) - : CudaKernel(info), epsilon_(info.GetAttrOrDefault("epsilon", 1e-6f)) { - ORT_ENFORCE(epsilon_ >= 0.0f, "EmbeddingRMSNorm epsilon must be non-negative"); -} - -Status EmbeddingRMSNorm::ComputeInternal(OpKernelContext* context) const { - const Tensor* ids = context->Input(0); - const Tensor* table = context->Input(1); - const Tensor* weight = context->Input(2); - const Tensor* committed = context->InputCount() > 3 ? context->Input(3) : nullptr; - ORT_RETURN_IF_NOT(table->Shape().NumDimensions() == 2, - "EmbeddingRMSNorm table must have shape [vocabulary, hidden]"); - const int64_t vocab_size = table->Shape()[0]; - const int64_t hidden_size = table->Shape()[1]; - ORT_RETURN_IF_NOT(weight->Shape().NumDimensions() == 1 && - weight->Shape()[0] == hidden_size, - "EmbeddingRMSNorm weight must match the embedding hidden size"); - const int64_t rows = ids->Shape().Size(); - TensorShapeVector output_dims(ids->Shape().GetDims().begin(), ids->Shape().GetDims().end()); - output_dims.push_back(hidden_size); - Tensor* output = context->Output(0, TensorShape(output_dims)); - int64_t* slot_ids_ptr = nullptr; - int32_t* seqlens_k_ptr = nullptr; - int32_t* cu_seqlens_ptr = nullptr; - if (context->OutputCount() > 1) { - Tensor* slot_ids = context->Output(1, TensorShape({1, rows})); - if (slot_ids != nullptr) slot_ids_ptr = slot_ids->MutableData(); - } - if (context->OutputCount() > 2) { - Tensor* seqlens_k = context->Output(2, TensorShape({1})); - if (seqlens_k != nullptr) seqlens_k_ptr = seqlens_k->MutableData(); - } - if (context->OutputCount() > 3) { - Tensor* cu_seqlens = context->Output(3, TensorShape({2})); - if (cu_seqlens != nullptr) cu_seqlens_ptr = cu_seqlens->MutableData(); - } - const int32_t* committed_ptr = committed != nullptr ? committed->Data() : nullptr; - if (table->IsDataType()) { - const int8_t* table_i8 = table->Data(); - const float* row_scales = reinterpret_cast(table_i8 + vocab_size * hidden_size); - return LaunchEmbeddingRMSNormInt8(Stream(context), ids->Data(), table_i8, row_scales, - weight->Data(), epsilon_, - reinterpret_cast(output->MutableData()), - rows, vocab_size, hidden_size, - committed_ptr, slot_ids_ptr, seqlens_k_ptr, cu_seqlens_ptr); - } - return LaunchEmbeddingRMSNorm(Stream(context), ids->Data(), - reinterpret_cast(table->Data()), - weight->Data(), epsilon_, - reinterpret_cast(output->MutableData()), - rows, vocab_size, hidden_size, - committed_ptr, slot_ids_ptr, seqlens_k_ptr, cu_seqlens_ptr); -} - -Status RowSelect::ComputeInternal(OpKernelContext* context) const { - const Tensor* selection = context->Input(0); - ORT_RETURN_IF_NOT(selection->Shape().Size() == 1, "RowSelect selection must have one element"); - std::array inputs{}; - std::array outputs{}; - std::array row_sizes{}; - int64_t num_rows = 2; - for (int i = 0; i < 3; ++i) { - const Tensor* input = context->Input(i + 1); - const auto& shape = input->Shape(); - ORT_RETURN_IF_NOT(shape.NumDimensions() >= 2 && shape[0] == 1 && shape[1] >= 1, - "RowSelect inputs must have shape [1, R, ...]"); - num_rows = shape[1]; - row_sizes[i] = shape.Size() / shape[1]; - TensorShapeVector output_dims(shape.GetDims().begin(), shape.GetDims().end()); - output_dims[1] = 1; - Tensor* output = context->Output(i, TensorShape(output_dims)); - inputs[i] = reinterpret_cast(input->Data()); - outputs[i] = reinterpret_cast(output->MutableData()); - } - return LaunchRowSelect(Stream(context), selection->Data(), inputs.data(), - outputs.data(), row_sizes.data(), num_rows); -} - SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) : CudaKernel(info), head_dim_(narrow(info.GetAttrOrDefault("head_dim", 0))), @@ -267,22 +166,6 @@ Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { dense_prefill ? false : sigmoid_in_fp32_, silu_gate_); } -UntileHyperConnectionWeights::UntileHyperConnectionWeights(const OpKernelInfo& info) - : CudaKernel(info), kind_(static_cast(info.GetAttrOrDefault("kind", -1))) { - ORT_ENFORCE(kind_ == 0 || kind_ == 1, "UntileHyperConnectionWeights kind must be 0 or 1"); -} - -Status UntileHyperConnectionWeights::ComputeInternal(OpKernelContext* context) const { - const Tensor* input = context->Input(0); - ORT_RETURN_IF_NOT(input->Shape() == TensorShape({160, 20480}), - "UntileHyperConnectionWeights input must have shape [160, 20480]"); - const TensorShape output_shape = kind_ == 0 ? TensorShape({10240, 320}) : TensorShape({320, 10240}); - Tensor* output = context->Output(0, output_shape); - return LaunchUntileHyperConnectionWeights( - Stream(context), reinterpret_cast(input->Data()), - reinterpret_cast(output->MutableData()), kind_); -} - template class ArgMaxLastAxis; template class ArgMaxLastAxis; template class ArgMaxLastAxis; diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu index f5fa0f8592a19..5414a35fc9f0e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -83,123 +83,6 @@ Status LaunchTypedArgmax(cudaStream_t stream, const T* input, int64_t rows, return CUDA_CALL(cudaGetLastError()); } -__global__ void EmbedNormKernel(const int64_t* ids, const uint16_t* table, - const float* weight, float epsilon, - uint16_t* output, int64_t vocab_size, int64_t hidden_size, - const int32_t* committed, int64_t* slot_ids, - int32_t* seqlens_k, int32_t* cu_seqlens, int64_t rows) { - __shared__ float sums[512]; - if (threadIdx.x == 0 && committed != nullptr) { - const int32_t c = committed[0]; - if (slot_ids != nullptr) slot_ids[blockIdx.x] = static_cast(c) + static_cast(blockIdx.x); - if (blockIdx.x == 0) { - if (seqlens_k != nullptr) seqlens_k[0] = c + static_cast(rows) - 1; - if (cu_seqlens != nullptr) { - cu_seqlens[0] = 0; - cu_seqlens[1] = static_cast(rows); - } - } - } - const int64_t row_index = blockIdx.x; - int64_t token_id = ids[row_index]; - if (token_id < -vocab_size || token_id >= vocab_size) { - for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { - output[row_index * hidden_size + col] = __bfloat16_as_ushort(__float2bfloat16_rn(CUDART_NAN_F)); - } - return; - } - if (token_id < 0) token_id += vocab_size; - const uint16_t* embedding = table + token_id * hidden_size; - float sum = 0.0f; - for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { - const float value = __uint_as_float(static_cast(embedding[col]) << 16); - sum += value * value; - } - sums[threadIdx.x] = sum; - __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - sums[threadIdx.x] += sums[threadIdx.x + stride]; - } - __syncthreads(); - } - const float scale = 1.0f / sqrtf(sums[0] / static_cast(hidden_size) + epsilon); - for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { - const float value = __uint_as_float(static_cast(embedding[col]) << 16); - const float normalized = (value * scale) * weight[col]; - output[row_index * hidden_size + col] = - __bfloat16_as_ushort(__float2bfloat16_rn(normalized)); - } -} - -__global__ void EmbedNormKernelInt8(const int64_t* ids, const int8_t* table_i8, const float* row_scales, - const float* weight, float epsilon, - uint16_t* output, int64_t vocab_size, int64_t hidden_size, - const int32_t* committed, int64_t* slot_ids, - int32_t* seqlens_k, int32_t* cu_seqlens, int64_t rows) { - __shared__ float sums[512]; - if (threadIdx.x == 0 && committed != nullptr) { - const int32_t c = committed[0]; - if (slot_ids != nullptr) slot_ids[blockIdx.x] = static_cast(c) + static_cast(blockIdx.x); - if (blockIdx.x == 0) { - if (seqlens_k != nullptr) seqlens_k[0] = c + static_cast(rows) - 1; - if (cu_seqlens != nullptr) { - cu_seqlens[0] = 0; - cu_seqlens[1] = static_cast(rows); - } - } - } - const int64_t row_index = blockIdx.x; - int64_t token_id = ids[row_index]; - if (token_id < -vocab_size || token_id >= vocab_size) { - for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { - output[row_index * hidden_size + col] = __bfloat16_as_ushort(__float2bfloat16_rn(CUDART_NAN_F)); - } - return; - } - if (token_id < 0) token_id += vocab_size; - const int8_t* embedding = table_i8 + token_id * hidden_size; - const float rscale = row_scales[token_id]; - float sum = 0.0f; - for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { - const float value = __bfloat162float(__float2bfloat16_rn(static_cast(embedding[col]) * rscale)); - sum += value * value; - } - sums[threadIdx.x] = sum; - __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - sums[threadIdx.x] += sums[threadIdx.x + stride]; - } - __syncthreads(); - } - const float scale = 1.0f / sqrtf(sums[0] / static_cast(hidden_size) + epsilon); - for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { - const float value = __bfloat162float(__float2bfloat16_rn(static_cast(embedding[col]) * rscale)); - const float normalized = (value * scale) * weight[col]; - output[row_index * hidden_size + col] = - __bfloat16_as_ushort(__float2bfloat16_rn(normalized)); - } -} - -__global__ void RowSelectKernel(const int64_t* selection, - const uint16_t* input0, const uint16_t* input1, - const uint16_t* input2, uint16_t* output0, - uint16_t* output1, uint16_t* output2, - int64_t row_size0, int64_t row_size1, - int64_t row_size2, int64_t num_rows) { - const int which = blockIdx.y; - const uint16_t* input = which == 0 ? input0 : (which == 1 ? input1 : input2); - uint16_t* output = which == 0 ? output0 : (which == 1 ? output1 : output2); - const int64_t row_size = which == 0 ? row_size0 : (which == 1 ? row_size1 : row_size2); - const int64_t raw_sel = selection[0]; - const int64_t selected_row = raw_sel <= 0 ? 0 : (raw_sel >= num_rows ? (num_rows - 1) : raw_sel); - for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - i < row_size; i += static_cast(gridDim.x) * blockDim.x) { - output[i] = input[selected_row * row_size + i]; - } -} - __device__ __forceinline__ float StableSigmoid(float value) { if (value >= 0.0f) { return 1.0f / (1.0f + expf(-value)); @@ -215,12 +98,12 @@ __device__ __forceinline__ float RoundedBf16Sigmoid(float value, bool sigmoid_in using namespace onnxruntime::cuda; const BFloat16 a(value); return static_cast(a > BFloat16(0.0f) - ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) - : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a))); + ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) + : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a))); } __device__ float ProjectedRowGate(const __nv_bfloat16* activation, const __nv_bfloat16* weight, - int64_t row, int64_t input_width, bool sigmoid_in_fp32) { + int64_t row, int64_t input_width, bool sigmoid_in_fp32) { using Reduce = cub::BlockReduce; __shared__ typename Reduce::TempStorage reduction; __shared__ float gate; @@ -259,8 +142,40 @@ __global__ void ProjectedSigmoidGateMulKernel( int64_t input_width, int64_t output_width, bool sigmoid_in_fp32) { const int64_t row = blockIdx.x; const float gate = ProjectedRowGate(activation, weight, row, input_width, sigmoid_in_fp32); + const int64_t row_base = row * output_width; + if ((output_width & 7) == 0 && + (reinterpret_cast(attention + row_base) & 15) == 0 && + (reinterpret_cast(output + row_base) & 15) == 0 && + (residual == nullptr || (reinterpret_cast(residual + row_base) & 15) == 0)) { + const int64_t vec_width = output_width >> 3; + const uint4* attn_vec = reinterpret_cast(attention + row_base); + const uint4* res_vec = residual != nullptr ? reinterpret_cast(residual + row_base) : nullptr; + uint4* out_vec = reinterpret_cast(output + row_base); + for (int64_t v = threadIdx.x; v < vec_width; v += kThreads) { + const uint4 a_u4 = attn_vec[v]; + const __nv_bfloat16* a_bf = reinterpret_cast(&a_u4); + uint4 o_u4; + __nv_bfloat16* o_bf = reinterpret_cast<__nv_bfloat16*>(&o_u4); + if (res_vec != nullptr) { + const uint4 r_u4 = res_vec[v]; + const __nv_bfloat16* r_bf = reinterpret_cast(&r_u4); +#pragma unroll + for (int m = 0; m < 8; ++m) { + const __nv_bfloat16 scaled = __float2bfloat16_rn(__bfloat162float(a_bf[m]) * gate); + o_bf[m] = __float2bfloat16_rn(__bfloat162float(scaled) + __bfloat162float(r_bf[m])); + } + } else { +#pragma unroll + for (int m = 0; m < 8; ++m) { + o_bf[m] = __float2bfloat16_rn(__bfloat162float(a_bf[m]) * gate); + } + } + out_vec[v] = o_u4; + } + return; + } for (int64_t column = threadIdx.x; column < output_width; column += kThreads) { - const int64_t index = row * output_width + column; + const int64_t index = row_base + column; __nv_bfloat16 value = __float2bfloat16_rn(__bfloat162float(attention[index]) * gate); if (residual != nullptr) { value = __float2bfloat16_rn(__bfloat162float(value) + __bfloat162float(residual[index])); @@ -316,34 +231,6 @@ __global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention, output[i] = value; } -__global__ void UntileHcBf16Kernel(const uint16_t* input, uint16_t* output, int kind) { - constexpr int64_t kVectorCount = 204800; - constexpr int64_t kElementCount = kVectorCount * 16; - const int64_t element = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (element >= kElementCount) return; - const int64_t vector_index = element / 16; - const int64_t lane = element % 16; - int64_t destination_vector; - if (kind == 0) { - const int64_t rp = vector_index & 127; - const int64_t t1 = vector_index >> 7; - const int64_t i = t1 % 10; - const int64_t t2 = t1 / 10; - const int64_t j = t2 % 20; - const int64_t s = t2 / 20; - destination_vector = (((s * 10 + i) * 128 + rp) * 20 + j); - } else { - const int64_t c = vector_index & 3; - const int64_t t1 = vector_index >> 2; - const int64_t rp = t1 & 31; - const int64_t t2 = t1 >> 5; - const int64_t i = t2 % 10; - const int64_t hb = t2 / 10; - destination_vector = (((i * 32 + rp) * 4 + c) * 160 + hb); - } - output[destination_vector * 16 + lane] = input[element]; -} - } // namespace Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_type, @@ -367,52 +254,10 @@ Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_ty } } -Status LaunchEmbeddingRMSNorm(cudaStream_t stream, const int64_t* ids, const uint16_t* table, - const float* weight, float epsilon, uint16_t* output, - int64_t rows, int64_t vocab_size, int64_t hidden_size, - const int32_t* committed, int64_t* slot_ids, - int32_t* seqlens_k, int32_t* cu_seqlens) { - ORT_RETURN_IF_NOT(hidden_size > 0 && vocab_size > 0, - "EmbeddingRMSNorm requires non-empty table dimensions"); - ORT_RETURN_IF_NOT(rows <= std::numeric_limits::max(), - "EmbeddingRMSNorm has too many rows for a CUDA grid"); - if (rows == 0) return Status::OK(); - EmbedNormKernel<<(rows), 512, 0, stream>>>( - ids, table, weight, epsilon, output, vocab_size, hidden_size, - committed, slot_ids, seqlens_k, cu_seqlens, rows); - return CUDA_CALL(cudaGetLastError()); -} - -Status LaunchEmbeddingRMSNormInt8(cudaStream_t stream, const int64_t* ids, const int8_t* table_i8, - const float* row_scales, const float* weight, float epsilon, - uint16_t* output, int64_t rows, int64_t vocab_size, int64_t hidden_size, - const int32_t* committed, int64_t* slot_ids, - int32_t* seqlens_k, int32_t* cu_seqlens) { - ORT_RETURN_IF_NOT(hidden_size > 0 && vocab_size > 0, - "EmbeddingRMSNorm requires non-empty table dimensions"); - ORT_RETURN_IF_NOT(rows <= std::numeric_limits::max(), - "EmbeddingRMSNorm has too many rows for a CUDA grid"); - if (rows == 0) return Status::OK(); - EmbedNormKernelInt8<<(rows), 512, 0, stream>>>( - ids, table_i8, row_scales, weight, epsilon, output, vocab_size, hidden_size, - committed, slot_ids, seqlens_k, cu_seqlens, rows); - return CUDA_CALL(cudaGetLastError()); -} - -Status LaunchRowSelect(cudaStream_t stream, const int64_t* selection, - const uint16_t* const* inputs, uint16_t* const* outputs, - const int64_t* row_sizes, int64_t num_rows) { - const dim3 grid(8, 3); - RowSelectKernel<<>>( - selection, inputs[0], inputs[1], inputs[2], outputs[0], outputs[1], outputs[2], - row_sizes[0], row_sizes[1], row_sizes[2], num_rows); - return CUDA_CALL(cudaGetLastError()); -} - Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, - const uint16_t* activation, const uint16_t* weight, - uint16_t* prepared, float* row_gates, - int rows, int width, int input_width) { + const uint16_t* activation, const uint16_t* weight, + uint16_t* prepared, float* row_gates, + int rows, int width, int input_width) { ORT_RETURN_IF_NOT(rows >= 0 && width > 0 && input_width > 0, "Shared expert preparation requires nonnegative rows and positive widths"); if (rows == 0) return Status::OK(); @@ -478,13 +323,4 @@ Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, return CUDA_CALL(cudaGetLastError()); } -Status LaunchUntileHyperConnectionWeights(cudaStream_t stream, const uint16_t* input, - uint16_t* output, int kind) { - constexpr int64_t kElementCount = 204800ll * 16; - constexpr int threads = 256; - constexpr int blocks = static_cast((kElementCount + threads - 1) / threads); - UntileHcBf16Kernel<<>>(input, output, kind); - return CUDA_CALL(cudaGetLastError()); -} - } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h index 88ebac79bb6ee..b0e7593371cd8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h @@ -15,19 +15,6 @@ Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_ty inline int64_t GetArgMaxLastAxisPartitions(int64_t width) { return 1 + (width - 1) / 4096; } -Status LaunchEmbeddingRMSNorm(cudaStream_t stream, const int64_t* ids, const uint16_t* table, - const float* weight, float epsilon, uint16_t* output, - int64_t rows, int64_t vocab_size, int64_t hidden_size, - const int32_t* committed = nullptr, int64_t* slot_ids = nullptr, - int32_t* seqlens_k = nullptr, int32_t* cu_seqlens = nullptr); -Status LaunchEmbeddingRMSNormInt8(cudaStream_t stream, const int64_t* ids, const int8_t* table_i8, - const float* row_scales, const float* weight, float epsilon, - uint16_t* output, int64_t rows, int64_t vocab_size, int64_t hidden_size, - const int32_t* committed = nullptr, int64_t* slot_ids = nullptr, - int32_t* seqlens_k = nullptr, int32_t* cu_seqlens = nullptr); -Status LaunchRowSelect(cudaStream_t stream, const int64_t* selection, - const uint16_t* const* inputs, uint16_t* const* outputs, - const int64_t* row_sizes, int64_t num_rows = 2); Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, const uint16_t* projection, uint16_t* output, int64_t count, int head_dim, int heads, int q_stride, @@ -40,11 +27,9 @@ Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attent int64_t rows, int64_t input_width, int64_t output_width, bool sigmoid_in_fp32); Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, - const uint16_t* activation, const uint16_t* weight, - uint16_t* prepared, float* row_gates, - int rows, int width, int input_width); -Status LaunchUntileHyperConnectionWeights(cudaStream_t stream, const uint16_t* input, - uint16_t* output, int kind); + const uint16_t* activation, const uint16_t* weight, + uint16_t* prepared, float* row_gates, + int rows, int width, int input_width); template class ArgMaxLastAxis final : public onnxruntime::cuda::CudaKernel { @@ -56,21 +41,6 @@ class ArgMaxLastAxis final : public onnxruntime::cuda::CudaKernel { bool keepdims_; }; -class EmbeddingRMSNorm final : public onnxruntime::cuda::CudaKernel { - public: - explicit EmbeddingRMSNorm(const OpKernelInfo& info); - Status ComputeInternal(OpKernelContext* context) const override; - - private: - float epsilon_; -}; - -class RowSelect final : public onnxruntime::cuda::CudaKernel { - public: - explicit RowSelect(const OpKernelInfo& info) : CudaKernel(info) {} - Status ComputeInternal(OpKernelContext* context) const override; -}; - class SigmoidGateMul final : public onnxruntime::cuda::CudaKernel { public: explicit SigmoidGateMul(const OpKernelInfo& info); @@ -87,13 +57,4 @@ class SigmoidGateMul final : public onnxruntime::cuda::CudaKernel { bool silu_gate_; }; -class UntileHyperConnectionWeights final : public onnxruntime::cuda::CudaKernel { - public: - explicit UntileHyperConnectionWeights(const OpKernelInfo& info); - Status ComputeInternal(OpKernelContext* context) const override; - - private: - int kind_; -}; - } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 9e0c203e834c4..eff8e87712966 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -173,10 +173,7 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, GatedDeltaNet); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, ArgMaxLastAxis); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, ArgMaxLastAxis); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, ArgMaxLastAxis); -class CUDA_MS_OP_CLASS_NAME(1, EmbeddingRMSNorm); -class CUDA_MS_OP_CLASS_NAME(1, RowSelect); class CUDA_MS_OP_CLASS_NAME(1, SigmoidGateMul); -class CUDA_MS_OP_CLASS_NAME(1, UntileHyperConnectionWeights); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, LinearAttentionGate); @@ -544,10 +541,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index cef02a698d5c9..e07a3dcf83f60 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4616,17 +4616,16 @@ ONNX_MS_OPERATOR_SET_SCHEMA( 1.0e-5f) .Input(0, "key", - "Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size) " - "or (total_tokens, hc_mult, hidden_size).", + "Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size).", "T") .Input(1, "query", - "Hidden-state queries with the same shape as key.", + "Hidden-state queries with shape (batch_size, sequence_length, hc_mult, hidden_size).", "T") .Input(2, "value", "Projected Engram value shared by every hyper-connection, with shape " - "(batch_size, sequence_length, hidden_size) or (total_tokens, hidden_size).", + "(batch_size, sequence_length, hidden_size).", "T") .Input(3, "key_norm_scale", @@ -4644,12 +4643,12 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Output(0, "output", - "Gated value tensor with the same shape as key.", + "Gated value tensor with shape (batch_size, sequence_length, hc_mult, hidden_size).", "T") .Output(1, "gated_value_normed", "Optional RMS-normalized gated value tensor with shape " - "(batch_size, sequence_length, hc_mult, hidden_size) or (total_tokens, hc_mult, hidden_size).", + "(batch_size, sequence_length, hc_mult, hidden_size).", "T", OpSchema::Optional) .TypeConstraint("T", @@ -4663,8 +4662,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( if (hasInputShape(ctx, 0)) { const auto& key_shape = getInputShape(ctx, 0); - if (key_shape.dim_size() != 3 && key_shape.dim_size() != 4) { - fail_shape_inference("EngramGate: key must have rank 3 or 4"); + if (key_shape.dim_size() != 4) { + fail_shape_inference("EngramGate: key must have rank 4"); } propagateShapeFromInputToOutput(ctx, 0, 0); if (ctx.getNumOutputs() > 1) { @@ -4673,21 +4672,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } if (hasInputShape(ctx, 1)) { const auto& query_shape = getInputShape(ctx, 1); - if (query_shape.dim_size() != 3 && query_shape.dim_size() != 4) { - fail_shape_inference("EngramGate: query must have rank 3 or 4"); - } - if (hasInputShape(ctx, 0) && query_shape.dim_size() != getInputShape(ctx, 0).dim_size()) { - fail_shape_inference("EngramGate: query rank must match key rank"); + if (query_shape.dim_size() != 4) { + fail_shape_inference("EngramGate: query must have rank 4"); } } if (hasInputShape(ctx, 2)) { const auto& value_shape = getInputShape(ctx, 2); - if (value_shape.dim_size() != 2 && value_shape.dim_size() != 3) { - fail_shape_inference("EngramGate: value must have rank 2 or 3"); - } - if ((hasInputShape(ctx, 0) && value_shape.dim_size() + 1 != getInputShape(ctx, 0).dim_size()) || - (hasInputShape(ctx, 1) && value_shape.dim_size() + 1 != getInputShape(ctx, 1).dim_size())) { - fail_shape_inference("EngramGate: value rank must be one lower than key and query rank"); + if (value_shape.dim_size() != 3) { + fail_shape_inference("EngramGate: value must have rank 3"); } } })); @@ -5762,75 +5754,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( updateOutputShape(ctx, 0, output_shape); })); -ONNX_MS_OPERATOR_SET_SCHEMA( - EmbeddingRMSNorm, 1, - OpSchema() - .SetDoc(FusedTensorOps_ver1_doc) - .Attr("epsilon", "RMS normalization epsilon.", AttributeProto::FLOAT) - .Input(0, "ids", - "Embedding row indices with any shape. Negative indices in [-vocabulary, -1] " - "count from the end. Out-of-range indices produce NaN output rows.", - "I") - .Input(1, "table", "BF16 or INT8 embedding matrix [vocabulary, hidden].", "TT") - .Input(2, "weight", "FP32 RMS normalization weight [hidden].", "F") - .Input(3, "committed_length", "Optional int32[1] committed length for head capture.", "I32", OpSchema::Optional) - .Output(0, "normalized", "BF16 normalized embeddings [shape(ids), hidden].", "T") - .Output(1, "slot_ids", "Optional int64[1, R] slot ids.", "I", OpSchema::Optional) - .Output(2, "seqlens_k", "Optional int32[1] seqlens_k.", "I32", OpSchema::Optional) - .Output(3, "cumulative_sequence_lengths", "Optional int32[2] cumulative sequence lengths.", "I32", OpSchema::Optional) - .TypeConstraint("I", {"tensor(int64)"}, "Indices use int64.") - .TypeConstraint("I32", {"tensor(int32)"}, "Committed length and sequence lengths use int32.") - .TypeConstraint("TT", {"tensor(bfloat16)", "tensor(int8)"}, "Embedding table uses BF16 or INT8.") - .TypeConstraint("T", {"tensor(bfloat16)"}, "Normalized result uses BF16.") - .TypeConstraint("F", {"tensor(float)"}, "Normalization weight uses FP32.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - updateOutputElemType(ctx, 0, ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); - if (ctx.getNumOutputs() > 1 && ctx.hasOutput(1)) { - updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_INT64); - } - if (ctx.getNumOutputs() > 2 && ctx.hasOutput(2)) { - updateOutputElemType(ctx, 2, ONNX_NAMESPACE::TensorProto_DataType_INT32); - } - if (ctx.getNumOutputs() > 3 && ctx.hasOutput(3)) { - updateOutputElemType(ctx, 3, ONNX_NAMESPACE::TensorProto_DataType_INT32); - } - if (!hasInputShape(ctx, 0) || !hasInputShape(ctx, 1)) return; - const auto& ids_shape = getInputShape(ctx, 0); - const auto& table_shape = getInputShape(ctx, 1); - if (table_shape.dim_size() != 2) { - fail_shape_inference("EmbeddingRMSNorm: table must have rank two"); - } - ONNX_NAMESPACE::TensorShapeProto output_shape = ids_shape; - *output_shape.add_dim() = table_shape.dim(1); - updateOutputShape(ctx, 0, output_shape); - })); - -ONNX_MS_OPERATOR_SET_SCHEMA( - RowSelect, 1, - OpSchema() - .SetDoc(FusedTensorOps_ver1_doc) - .Input(0, "selection", "Scalar-like int64 row selector; zero selects row 0.", "I") - .Input(1, "x0", "BF16 tensor with shape [1, 2, ...].", "T") - .Input(2, "x1", "BF16 tensor with shape [1, 2, ...].", "T") - .Input(3, "x2", "BF16 tensor with shape [1, 2, ...].", "T") - .Output(0, "y0", "Selected BF16 tensor with shape [1, 1, ...].", "T") - .Output(1, "y1", "Selected BF16 tensor with shape [1, 1, ...].", "T") - .Output(2, "y2", "Selected BF16 tensor with shape [1, 1, ...].", "T") - .TypeConstraint("I", {"tensor(int64)"}, "Selector uses int64.") - .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and outputs use BF16.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - for (int i = 0; i < 3; ++i) { - propagateElemTypeFromInputToOutput(ctx, i + 1, i); - if (!hasInputShape(ctx, i + 1)) continue; - ONNX_NAMESPACE::TensorShapeProto output_shape = getInputShape(ctx, i + 1); - if (output_shape.dim_size() < 2) { - fail_shape_inference("RowSelect: inputs must have rank at least two"); - } - output_shape.mutable_dim(1)->set_dim_value(1); - updateOutputShape(ctx, i, output_shape); - } - })); - ONNX_MS_OPERATOR_SET_SCHEMA( SigmoidGateMul, 1, OpSchema() @@ -5869,29 +5792,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( if (hasInputShape(ctx, 0)) updateOutputShape(ctx, 0, getInputShape(ctx, 0)); })); -ONNX_MS_OPERATOR_SET_SCHEMA( - UntileHyperConnectionWeights, 1, - OpSchema() - .SetDoc(FusedTensorOps_ver1_doc) - .Attr("kind", "Untile layout kind, either 0 or 1.", AttributeProto::INT) - .Input(0, "tiled", "Tiled BF16 weight [160, 20480].", "T") - .Input(1, "unused", "Reserved optional input retained for exporter compatibility.", "T", - OpSchema::Optional) - .Output(0, "untiled", "Untiled BF16 weight.", "T") - .TypeConstraint("T", {"tensor(bfloat16)"}, "Input and output use BF16.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - if (!hasInputShape(ctx, 0)) return; - const int64_t kind = getAttribute(ctx, "kind", static_cast(-1)); - if (kind != 0 && kind != 1) { - fail_shape_inference("UntileHyperConnectionWeights: kind must be zero or one"); - } - ONNX_NAMESPACE::TensorShapeProto output_shape; - output_shape.add_dim()->set_dim_value(kind == 0 ? 10240 : 320); - output_shape.add_dim()->set_dim_value(kind == 0 ? 320 : 10240); - updateOutputShape(ctx, 0, output_shape); - })); - constexpr const char* LinearAttentionGate_ver1_doc = R"DOC( Fuses the gate projections that feed LinearAttention's gated-delta recurrence: diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 523f9d4a6a4c3..7ae6abc106f12 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -97,10 +97,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPostMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNet); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNetStateReplay); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ArgMaxLastAxis); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, EmbeddingRMSNorm); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RowSelect); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SigmoidGateMul); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, UntileHyperConnectionWeights); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ConvStateReplay); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedRMSNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedAdd); @@ -237,10 +234,7 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); - fn(GetOpSchema()); - fn(GetOpSchema()); fn(GetOpSchema()); - fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 605c32578c386..2af2fd6b49e6c 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -783,14 +783,6 @@ std::vector> WebGpuExecutionProvider::GetCapa } } - if (node.OpType() == "EngramGate" && node.Domain() == kMSDomain) { - const auto* key_shape = node.InputDefs()[0]->Shape(); - if (key_shape == nullptr || key_shape->dim_size() != 4) { - LOGS(*GetLogger(), INFO) << "EngramGate requires rank-4 key input for WebGPU; falling back to CPU"; - continue; - } - } - // Check for MatMulBnb4 if (node.OpType() == "MatMulBnb4" && node.Domain() == kMSDomain) { // Current implementation only supports the forward case (transB=1). transB defaults to 1. diff --git a/onnxruntime/test/contrib_ops/engram_ops_test.cc b/onnxruntime/test/contrib_ops/engram_ops_test.cc index 66d3c0a8085f1..0471be6941f53 100644 --- a/onnxruntime/test/contrib_ops/engram_ops_test.cc +++ b/onnxruntime/test/contrib_ops/engram_ops_test.cc @@ -9,7 +9,6 @@ #include #include #include -#include #include #include #include @@ -70,16 +69,13 @@ bool IsTypeSupported() { // Runs the tester on CUDA only for BFloat16, otherwise on the default set of execution providers. template -void RunOnSupportedProviders(OpTester& test, bool packed = false) { +void RunOnSupportedProviders(OpTester& test) { if constexpr (std::is_same_v) { std::vector> providers; providers.push_back(DefaultCudaExecutionProvider()); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } else { - // OpTester preassigns nodes without GetCapability; the session regression below tests fallback. - test.Run(OpTester::ExpectResult::kExpectSuccess, "", - packed ? std::unordered_set{kWebGpuExecutionProvider} - : std::unordered_set{}); + test.Run(); } } @@ -155,7 +151,7 @@ void RunEngramGateTest(float tolerance) { // Verifies gated_value_normed: RMSNorm is applied independently to each hyper-connection branch // (each hidden_size-sized slice), not over the concatenated hc_mult * hidden_size dimension. template -void RunEngramGateNormedTest(float tolerance, bool packed = false) { +void RunEngramGateNormedTest(float tolerance) { if (!IsTypeSupported()) { GTEST_SKIP() << "No execution provider available for this type"; } @@ -208,21 +204,17 @@ void RunEngramGateNormedTest(float tolerance, bool packed = false) { OpTester test("EngramGate", 1, kMSDomain); test.AddAttribute("epsilon", kEpsilon); - const std::vector key_shape = packed ? std::vector{1, hc_mult, hidden_size} - : std::vector{1, 1, hc_mult, hidden_size}; - const std::vector value_shape = packed ? std::vector{1, hidden_size} - : std::vector{1, 1, hidden_size}; - test.AddInput("key", key_shape, ToTensorType(key)); - test.AddInput("query", key_shape, ToTensorType(query)); - test.AddInput("value", value_shape, ToTensorType(value)); + test.AddInput("key", {1, 1, hc_mult, hidden_size}, ToTensorType(key)); + test.AddInput("query", {1, 1, hc_mult, hidden_size}, ToTensorType(query)); + test.AddInput("value", {1, 1, hidden_size}, ToTensorType(value)); test.AddInput("key_norm_scale", {hc_mult, hidden_size}, ToTensorType(key_scale)); test.AddInput("query_norm_scale", {hc_mult, hidden_size}, ToTensorType(query_scale)); test.AddInput("conv_norm_scale", {hc_mult, hidden_size}, ToTensorType(conv_scale)); - test.AddOutput("output", key_shape, ToTensorType(gated_value), false, tolerance, + test.AddOutput("output", {1, 1, hc_mult, hidden_size}, ToTensorType(gated_value), false, tolerance, tolerance); - test.AddOutput("gated_value_normed", key_shape, ToTensorType(expected_normed), false, + test.AddOutput("gated_value_normed", {1, 1, hc_mult, hidden_size}, ToTensorType(expected_normed), false, tolerance, tolerance); - RunOnSupportedProviders(test, packed); + RunOnSupportedProviders(test); } // Exercises hc_mult > 1 and non-unit norm scales for an arbitrary hidden_size. hidden_size == 4 @@ -233,7 +225,7 @@ void RunEngramGateNormedTest(float tolerance, bool packed = false) { // 64 over hidden_size / components), so only a hidden_size above those strides takes a second // iteration and actually accumulates into the per-thread partials. template -void RunEngramGateVectorizedTest(float tolerance, int64_t hidden, bool packed = false) { +void RunEngramGateVectorizedTest(float tolerance, int64_t hidden) { if (!IsTypeSupported()) { GTEST_SKIP() << "No execution provider available for this type"; } @@ -279,18 +271,14 @@ void RunEngramGateVectorizedTest(float tolerance, int64_t hidden, bool packed = OpTester test("EngramGate", 1, kMSDomain); test.AddAttribute("epsilon", kEpsilon); - const std::vector key_shape = packed ? std::vector{kBatch * kSequence, kHcMult, hidden} - : std::vector{kBatch, kSequence, kHcMult, hidden}; - const std::vector value_shape = packed ? std::vector{kBatch * kSequence, hidden} - : std::vector{kBatch, kSequence, hidden}; - test.AddInput("key", key_shape, ToTensorType(key)); - test.AddInput("query", key_shape, ToTensorType(query)); - test.AddInput("value", value_shape, ToTensorType(value)); + test.AddInput("key", {kBatch, kSequence, kHcMult, hidden}, ToTensorType(key)); + test.AddInput("query", {kBatch, kSequence, kHcMult, hidden}, ToTensorType(query)); + test.AddInput("value", {kBatch, kSequence, hidden}, ToTensorType(value)); test.AddInput("key_norm_scale", {kHcMult, hidden}, ToTensorType(key_scale)); test.AddInput("query_norm_scale", {kHcMult, hidden}, ToTensorType(query_scale)); - test.AddOutput("output", key_shape, ToTensorType(expected), false, tolerance, + test.AddOutput("output", {kBatch, kSequence, kHcMult, hidden}, ToTensorType(expected), false, tolerance, tolerance); - RunOnSupportedProviders(test, packed); + RunOnSupportedProviders(test); } // --------------------------------------------------------------------------------------------- @@ -1805,109 +1793,6 @@ TEST(EngramOpsTest, EngramGateVectorizedFloat16) { RunEngramGateVectorizedTest(3e-3f, 4); } -TEST(EngramOpsTest, EngramGatePackedFloat) { - RunEngramGateVectorizedTest(1e-4f, 4, true); -} - -TEST(EngramOpsTest, EngramGatePackedFloat16) { - RunEngramGateVectorizedTest(3e-3f, 4, true); -} - -TEST(EngramOpsTest, EngramGatePackedBFloat16) { - RunEngramGateVectorizedTest(2e-2f, 4, true); -} - -#ifdef USE_WEBGPU -template -void RunEngramGatePackedWebGpuCpuFallbackTest() { - for (bool normed : {false, true}) { - auto webgpu = DefaultWebGpuExecutionProvider(); - if (webgpu == nullptr) GTEST_SKIP() << "WebGPU execution provider is not available"; - std::unordered_map domain_to_version{{kOnnxDomain, 17}, {kMSDomain, 1}}; - Model model("EngramGatePackedFallback", false, ModelMetaData(), PathString(), - IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, {}, - DefaultLoggingManager().DefaultLogger()); - Graph& graph = model.MainGraph(); - ModelTestBuilder builder(graph); - auto* key = builder.MakeInput({2, 2, 1}, ToTensorType({1, 1, 1, 1})); - auto* query = builder.MakeInput({2, 2, 1}, ToTensorType({1, 1, 1, 1})); - auto* value = builder.MakeInput({2, 1}, ToTensorType({2, 4})); - auto* key_scale = builder.MakeInput({2, 1}, ToTensorType({1, 1})); - auto* query_scale = builder.MakeInput({2, 1}, ToTensorType({1, 1})); - std::vector inputs{key, query, value, key_scale, query_scale}; - std::vector outputs{builder.MakeOutput()}; - if (normed) { - inputs.push_back(builder.MakeInput({2, 1}, ToTensorType({1, 1}))); - outputs.push_back(builder.MakeOutput()); - } - auto& node = builder.AddNode("EngramGate", inputs, outputs, kMSDomain); - node.AddAttribute("epsilon", 0.0f); - builder.SetGraphOutputs(); - ASSERT_STATUS_OK(graph.Resolve()); - ONNX_NAMESPACE::ModelProto model_proto = model.ToProto(); - std::string model_data; - ASSERT_TRUE(model_proto.SerializeToString(&model_data)); - SessionOptions options; - options.graph_optimization_level = TransformerLevel::Default; - InferenceSessionWrapper session(options, GetEnvironment()); - ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(webgpu))); - ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultCpuExecutionProvider())); - std::istringstream model_stream(model_data); - ASSERT_STATUS_OK(session.Load(model_stream)); - ASSERT_STATUS_OK(session.Initialize()); - size_t gate_nodes = 0; - for (const auto& assigned_node : session.GetSessionState().GetGraphViewer().Nodes()) { - if (assigned_node.OpType() == "EngramGate") { - ++gate_nodes; - EXPECT_EQ(assigned_node.GetExecutionProviderType(), kCpuExecutionProvider); - } - } - ASSERT_EQ(gate_nodes, 1u); - std::vector fetches; - ASSERT_STATUS_OK(session.Run(RunOptions{}, builder.feeds_, builder.output_names_, &fetches)); - ASSERT_EQ(fetches.size(), normed ? 2u : 1u); - const float gate = Sigmoid(1.0f); - const auto expected = ToTensorType({2 * gate, 2 * gate, 4 * gate, 4 * gate}); - const auto actual = fetches[0].Get().DataAsSpan(); - ASSERT_EQ(actual.size(), expected.size()); - for (size_t i = 0; i < actual.size(); ++i) { - EXPECT_NEAR(static_cast(actual[i]), static_cast(expected[i]), 2e-3f); - } - if (normed) { - ASSERT_EQ(fetches[1].Get().Shape(), TensorShape({2, 2, 1})); - for (T element : fetches[1].Get().DataAsSpan()) { - EXPECT_NEAR(static_cast(element), 1.0f, 2e-3f); - } - } - } -} - -TEST(EngramOpsTest, EngramGatePackedWebGpuCpuFallbackFloat) { - RunEngramGatePackedWebGpuCpuFallbackTest(); -} - -TEST(EngramOpsTest, EngramGatePackedWebGpuCpuFallbackFloat16) { - RunEngramGatePackedWebGpuCpuFallbackTest(); -} -#endif - -TEST(EngramOpsTest, EngramGateRejectsMixedPackedRanks) { -#if defined(ORT_NO_EXCEPTIONS) - GTEST_SKIP() << "Static shape inference failures abort when exceptions are disabled."; -#endif - for (bool bad_query : {false, true}) { - OpTester test("EngramGate", 1, kMSDomain); - test.AddInput("key", {1, 1, 2}, {1, 1}); - test.AddInput("query", bad_query ? std::vector{1, 1, 1, 2} : std::vector{1, 1, 2}, {1, 1}); - test.AddInput("value", bad_query ? std::vector{1, 2} : std::vector{1, 1, 2}, {1, 1}); - test.AddInput("key_norm_scale", {1, 2}, {1, 1}); - test.AddInput("query_norm_scale", {1, 2}, {1, 1}); - test.AddOutput("output", {1, 1, 2}, {0, 0}); - test.Run(OpTester::ExpectResult::kExpectFailure, - bad_query ? "query rank must match key rank" : "value rank must be one lower"); - } -} - // 260 channels is the smallest multiple of 4 above both accumulation strides (CUDA's blockDim.x of // 256 and, with components == 4, the WGSL workgroup size of 64 over hidden_size / 4 == 65), so this // is the only case where either strided reduction loop runs more than once per thread. @@ -1931,14 +1816,6 @@ TEST(EngramOpsTest, EngramGateNormedFloat16) { RunEngramGateNormedTest(2e-3f); } -TEST(EngramOpsTest, EngramGatePackedNormedFloat) { - RunEngramGateNormedTest(1e-4f, true); -} - -TEST(EngramOpsTest, EngramGatePackedNormedBFloat16) { - RunEngramGateNormedTest(2e-2f, true); -} - TEST(EngramOpsTest, EngramGateNormedBFloat16) { RunEngramGateNormedTest(2e-2f); } diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc index 8b38aced87697..15901701d7ad9 100644 --- a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc +++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc @@ -97,59 +97,6 @@ TEST(FusedTensorOpsTest, ArgMaxLastAxisNaNSingleAndMultiplePartitions) { RunArgMaxNaNRows(); } -TEST(FusedTensorOpsTest, EmbeddingRMSNorm) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - - const std::vector table = {3.0f, 4.0f, 0.0f, 0.0f, 0.0f, 0.0f, 2.0f, 0.0f}; - const float inverse_first = 1.0f / std::sqrt(6.25f + 1e-6f); - const float inverse_second = 1.0f / std::sqrt(1.0f + 1e-6f); - const std::vector expected = { - 3.0f * inverse_first, 4.0f * inverse_first, 0.0f, 0.0f, - 0.0f, 0.0f, 2.0f * inverse_second, 0.0f}; - - OpTester test("EmbeddingRMSNorm", 1, kMSDomain); - test.AddAttribute("epsilon", 1e-6f); - test.AddInput("ids", {2}, {0, 1}); - test.AddInput("table", {2, 4}, ToTensorType(table)); - test.AddInput("weight", {4}, {1.0f, 1.0f, 1.0f, 1.0f}); - test.AddOutput("normalized", {2, 4}, ToTensorType(expected)); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - -TEST(FusedTensorOpsTest, EmbeddingRMSNormIndexBounds) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - const float nan = std::numeric_limits::quiet_NaN(); - OpTester test("EmbeddingRMSNorm", 1, kMSDomain); - test.AddAttribute("epsilon", 0.0f); - test.AddInput("ids", {8}, {0, 1, -1, -2, 2, -3, std::numeric_limits::min(), std::numeric_limits::max()}); - test.AddInput("table", {2, 4}, ToBFloat16({2, 0, 0, 0, 0, 2, 0, 0})); - test.AddInput("weight", {4}, {1, 1, 1, 1}); - test.AddOutput("normalized", {8, 4}, - ToBFloat16({2, 0, 0, 0, 0, 2, 0, 0, 0, 2, 0, 0, 2, 0, 0, 0, - nan, nan, nan, nan, nan, nan, nan, nan, - nan, nan, nan, nan, nan, nan, nan, nan})); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - -TEST(FusedTensorOpsTest, RowSelect) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - - OpTester test("RowSelect", 1, kMSDomain); - test.AddInput("selection", {}, {1}); - for (int i = 0; i < 3; ++i) { - const std::string input_name = "x" + std::to_string(i); - const std::string output_name = "y" + std::to_string(i); - test.AddInput(input_name.c_str(), {1, 2, 2}, - ToTensorType({1.0f, 2.0f, 3.0f, 4.0f})); - test.AddOutput(output_name.c_str(), {1, 1, 2}, - ToTensorType({3.0f, 4.0f})); - } - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - TEST(FusedTensorOpsTest, SigmoidGateMul) { auto providers = CudaProvider(); if (providers.empty()) GTEST_SKIP(); @@ -343,50 +290,5 @@ TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) { } } -TEST(FusedTensorOpsTest, UntileHyperConnectionWeights) { - constexpr int64_t kInputElements = 160ll * 20480; - constexpr int64_t kVectorCount = kInputElements / 16; - for (const int64_t kind : {0, 1}) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - OpTester test("UntileHyperConnectionWeights", 1, kMSDomain); - test.AddAttribute("kind", kind); - std::vector input(kInputElements); - std::vector expected(kInputElements); - for (int64_t i = 0; i < kInputElements; ++i) { - input[i] = BFloat16(static_cast(i % 127)); - } - for (int64_t vector_index = 0; vector_index < kVectorCount; ++vector_index) { - int64_t destination_vector; - if (kind == 0) { - const int64_t rp = vector_index & 127; - const int64_t t1 = vector_index >> 7; - const int64_t i = t1 % 10; - const int64_t t2 = t1 / 10; - const int64_t j = t2 % 20; - const int64_t s = t2 / 20; - destination_vector = (((s * 10 + i) * 128 + rp) * 20 + j); - } else { - const int64_t c = vector_index & 3; - const int64_t t1 = vector_index >> 2; - const int64_t rp = t1 & 31; - const int64_t t2 = t1 >> 5; - const int64_t i = t2 % 10; - const int64_t hb = t2 / 10; - destination_vector = (((i * 32 + rp) * 4 + c) * 160 + hb); - } - for (int64_t lane = 0; lane < 16; ++lane) { - expected[destination_vector * 16 + lane] = input[vector_index * 16 + lane]; - } - } - test.AddInput("tiled", {160, 20480}, input); - const std::vector output_shape = kind == 0 - ? std::vector{10240, 320} - : std::vector{320, 10240}; - test.AddOutput("untiled", output_shape, expected); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); - } -} - } // namespace } // namespace onnxruntime::test From 50dc7c024e424dde615caa294dc5b3105fe2671a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Oct 2026 23:41:06 +0000 Subject: [PATCH 11/19] Apply lintrunner formatting Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 12 ++++++------ onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu index 5414a35fc9f0e..8221df1706c61 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -98,12 +98,12 @@ __device__ __forceinline__ float RoundedBf16Sigmoid(float value, bool sigmoid_in using namespace onnxruntime::cuda; const BFloat16 a(value); return static_cast(a > BFloat16(0.0f) - ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) - : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a))); + ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) + : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a))); } __device__ float ProjectedRowGate(const __nv_bfloat16* activation, const __nv_bfloat16* weight, - int64_t row, int64_t input_width, bool sigmoid_in_fp32) { + int64_t row, int64_t input_width, bool sigmoid_in_fp32) { using Reduce = cub::BlockReduce; __shared__ typename Reduce::TempStorage reduction; __shared__ float gate; @@ -255,9 +255,9 @@ Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_ty } Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, - const uint16_t* activation, const uint16_t* weight, - uint16_t* prepared, float* row_gates, - int rows, int width, int input_width) { + const uint16_t* activation, const uint16_t* weight, + uint16_t* prepared, float* row_gates, + int rows, int width, int input_width) { ORT_RETURN_IF_NOT(rows >= 0 && width > 0 && input_width > 0, "Shared expert preparation requires nonnegative rows and positive widths"); if (rows == 0) return Status::OK(); diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h index b0e7593371cd8..4c57a089b886f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h @@ -27,9 +27,9 @@ Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attent int64_t rows, int64_t input_width, int64_t output_width, bool sigmoid_in_fp32); Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, - const uint16_t* activation, const uint16_t* weight, - uint16_t* prepared, float* row_gates, - int rows, int width, int input_width); + const uint16_t* activation, const uint16_t* weight, + uint16_t* prepared, float* row_gates, + int rows, int width, int input_width); template class ArgMaxLastAxis final : public onnxruntime::cuda::CudaKernel { From 6b8e19dd61a2037c96eb7f701befabb31c18aea4 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 04:09:35 +0000 Subject: [PATCH 12/19] fix(docs): sync SigmoidGateMul and ArgMaxLastAxis schema docs in ContribOperators.md Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 7d1a47b59839c..58a60d0f23ec0 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -158,7 +158,7 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
keepdims : int (default is 0)
+
keepdims : int
Keep a singleton final axis when nonzero.
@@ -7611,21 +7611,21 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
broadcast_gate : int (default is 0)
+
broadcast_gate : int
When one, projection has the attention shape with final dimension one. Packed layout attributes must be zero. CUDA only.
-
head_dim : int (default is 0)
+
head_dim : int
Dimension of each attention head.
-
num_heads : int (default is 0)
+
num_heads : int
Number of attention heads.
-
q_offset : int (default is 0)
+
q_offset : int
Offset of the query projection within each token.
-
q_stride : int (default is 0)
+
q_stride : int
Per-head stride in the packed projection.
-
sigmoid_in_fp32 : int (default is 1)
+
sigmoid_in_fp32 : int
One computes sigmoid in FP32 before BF16 conversion. Zero preserves the input-type intermediate arithmetic of CUDA Sigmoid, including subtraction on negative inputs.
-
silu_gate : int (default is 0)
+
silu_gate : int
When one, projection and attention have identical shapes. Multiply the BF16-rounded sigmoid by projection and round to BF16 before multiplying attention. Cannot be combined with broadcast_gate or packed layout attributes.
-
tok_stride : int (default is 0)
+
tok_stride : int
Per-token stride in the packed projection.
From ea860c3f5f4d8275ce2e8b8f7e37c7c1df88a3ed Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 25 Oct 2026 00:46:28 +0000 Subject: [PATCH 13/19] Replace ArgMaxLastAxis with BFloat16 ArgMax/ArgMin and extend GatedAdd with projected gate - Remove com.microsoft::ArgMaxLastAxis and instead register BFloat16 (opset 13+) on standard ai.onnx::ArgMax and ai.onnx::ArgMin in CUDA (reduction_ops.cc, reduction_functions.cu, cuda_execution_provider.cc). - Extend com.microsoft::GatedAdd with optional gate_weight [K, 1] projection, sigmoid_in_fp32 attribute, and 16-byte uint4 vectorization (GatedAddVec8Kernel), and vectorize ScaledSiLUKernel with uint4 (ScaledSiLUVec8Kernel). - Update fused_tensor_ops_test.cc, ContribOperators.md, and OperatorKernels.md. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 53 +++--------- docs/OperatorKernels.md | 9 +- .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 39 --------- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 86 ------------------- .../contrib_ops/cuda/bert/fused_tensor_ops.h | 16 ---- .../contrib_ops/cuda/bert/gated_add.cc | 23 ++++- onnxruntime/contrib_ops/cuda/bert/gated_add.h | 10 ++- .../contrib_ops/cuda/bert/gated_add_impl.cu | 41 +++++++++ .../contrib_ops/cuda/bert/scaled_silu_impl.cu | 31 +++++++ .../contrib_ops/cuda/cuda_contrib_kernels.cc | 6 -- .../core/graph/contrib_ops/bert_defs.cc | 37 ++------ onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 - .../providers/cuda/cuda_execution_provider.cc | 4 + .../cuda/reduction/reduction_functions.cu | 7 ++ .../providers/cuda/reduction/reduction_ops.cc | 14 +-- .../test/contrib_ops/fused_tensor_ops_test.cc | 18 ++-- 16 files changed, 156 insertions(+), 240 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 58a60d0f23ec0..bdcc34f44971e 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -3,7 +3,6 @@ Do not modify directly.* * com.microsoft - * com.microsoft.ArgMaxLastAxis * com.microsoft.Attention * com.microsoft.AttnLSTM * com.microsoft.BeamSearch @@ -147,45 +146,6 @@ Do not modify directly.* * experimental com.microsoft.QEmbedLayerNormalization ## com.microsoft -### **com.microsoft.ArgMaxLastAxis** - - CUDA contrib operators for fused tensor operations used by transformer models. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
keepdims : int
-
Keep a singleton final axis when nonzero.
-
- -#### Inputs - -
-
logits : T
-
Logits with vocabulary on the last axis.
-
- -#### Outputs - -
-
indices : I
-
First maximum index for each leading row. NaNs are ignored; an all-NaN row returns index zero.
-
- -#### Type Constraints - -
-
T : tensor(float), tensor(float16), tensor(bfloat16)
-
Supported logit types.
-
I : tensor(int64)
-
Argmax indices are int64.
-
- - ### **com.microsoft.Attention** Multi-Head Attention that can be either unidirectional (like GPT-2) or bidirectional (like BERT). @@ -2440,7 +2400,14 @@ This version of the operator has been available since version 1 of the 'com.micr This version of the operator has been available since version 1 of the 'com.microsoft' operator set. -#### Inputs +#### Attributes + +
+
sigmoid_in_fp32 : int
+
When gate_weight is provided, compute sigmoid in FP32 (1, default) or input type (0).
+
+ +#### Inputs (3 - 4)
X : T
@@ -2448,7 +2415,9 @@ This version of the operator has been available since version 1 of the 'com.micr
Y : T
Input scaled by gate, with the same shape as X.
gate : T
-
Per-row gate with shape (..., 1).
+
Per-row gate with shape (..., 1), or activation (..., K) when gate_weight is provided.
+
gate_weight (optional) : T
+
Optional projection weight [K, 1] applied to gate before sigmoid.
#### Outputs diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 42996435a79a1..a6628392f0837 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -600,7 +600,7 @@ The **OpSet Version** column uses the following notation: |FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |FusedGemm|*in* A:**T**
*in* B:**T**
*in* C:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float)| -|GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| +|GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*in* gate_weight:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |GatedRMSNorm|*in* X:**T**
*in* scale:**T**
*in* gate:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |GatherBlockQuantized|*in* data:**T1**
*in* indices:**Tind**
*in* scales:**T2**
*in* zero_points:**T1**
*out* output:**T2**|1+|**T1** = tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int4), tensor(uint4), tensor(uint8)
**T2** = tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)| |GatherND|*in* data:**T**
*in* indices:**Tind**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)| @@ -686,10 +686,10 @@ The **OpSet Version** column uses the following notation: |||[7, 12]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| |Affine|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |And|*in* A:**T**
*in* B:**T**
*out* C:**T1**|7+|**T** = tensor(bool)
**T1** = tensor(bool)| -|ArgMax|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(float16)| +|ArgMax|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |||12|**T** = tensor(double), tensor(float), tensor(float16)| |||[1, 11]|**T** = tensor(double), tensor(float), tensor(float16)| -|ArgMin|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(float16)| +|ArgMin|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |||12|**T** = tensor(double), tensor(float), tensor(float16)| |||[1, 11]|**T** = tensor(double), tensor(float), tensor(float16)| |Attention|*in* Q:**T1**
*in* K:**T1**
*in* V:**T2**
*in* attn_mask:**U**
*in* past_key:**T1**
*in* past_value:**T2**
*in* nonpad_kv_seqlen:**tensor(int64)**
*out* Y:**T1**
*out* present_key:**T1**
*out* present_value:**T2**
*out* qk_matmul_output:**T1**

or

*in* Q:**T1**
*in* K:**T1**
*in* V:**T2**
*in* attn_mask:**U**
*in* past_key:**T1**
*in* past_value:**T2**
*out* Y:**T1**
*out* present_key:**T1**
*out* present_value:**T2**
*out* qk_matmul_output:**T1**|24+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(bfloat16), tensor(float), tensor(float16)
**U** = tensor(bfloat16), tensor(bool), tensor(float), tensor(float16)| @@ -1093,7 +1093,6 @@ The **OpSet Version** column uses the following notation: | | | | |**Operator Domain:** *com.microsoft*|||| -|ArgMaxLastAxis|*in* logits:**T**
*out* indices:**I**|1+|**I** = tensor(int64)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |Attention|*in* input:**T**
*in* weights:**T**
*in* bias:**T**
*in* mask_index:**M**
*in* past:**T**
*in* attention_bias:**T**
*in* past_sequence_length:**M**
*out* output:**T**
*out* present:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |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), tensor(float16)| |BiasAdd|*in* X:**T**
*in* bias:**T**
*in* skip:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| @@ -1120,7 +1119,7 @@ The **OpSet Version** column uses the following notation: |FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| -|GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*in* gate_weight:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |GatedDeltaNet|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* cu_seqlens:**TI**
*in* decay:**TS**
*in* beta:**TS**
*in* initial_state:**TS**
*in* a_log:**TS**
*in* dt_bias:**TS**
*in* capture_count:**TI**
*in* state_update_active:**TI**
*out* output:**T**
*out* final_state:**TS**
*out* state_update:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)
**TS** = tensor(float)| |GatedRMSNorm|*in* X:**T**
*in* scale:**T**
*in* gate:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |GatedRelativePositionBias|*in* query_layer:**T**
*in* query_bias:**T**
*in* rel_pos:**T**
*in* weight:**T**
*in* bias:**T**
*in* eco_a:**T**
*in* token_offset:**M**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc index 1f360cf4d9b2a..ff1544dcf0c03 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -14,20 +14,6 @@ namespace onnxruntime::contrib::cuda { -#define REGISTER_ARGMAX_LAST_AXIS(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - ArgMaxLastAxis, kMSDomain, 1, T, kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ - ArgMaxLastAxis); - -REGISTER_ARGMAX_LAST_AXIS(float) -REGISTER_ARGMAX_LAST_AXIS(MLFloat16) -REGISTER_ARGMAX_LAST_AXIS(BFloat16) - -#undef REGISTER_ARGMAX_LAST_AXIS - ONNX_OPERATOR_KERNEL_EX( SigmoidGateMul, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) @@ -36,31 +22,6 @@ ONNX_OPERATOR_KERNEL_EX( .InputMemoryType(OrtMemTypeCPUInput, 3), SigmoidGateMul); -template -ArgMaxLastAxis::ArgMaxLastAxis(const OpKernelInfo& info) - : CudaKernel(info), keepdims_(info.GetAttrOrDefault("keepdims", 0) != 0) {} - -template -Status ArgMaxLastAxis::ComputeInternal(OpKernelContext* context) const { - const Tensor* input = context->Input(0); - const auto& shape = input->Shape(); - ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1 && - shape[shape.NumDimensions() - 1] > 0, - "ArgMaxLastAxis requires a non-empty final dimension"); - const int64_t width = shape[shape.NumDimensions() - 1]; - const int64_t rows = shape.Size() / width; - TensorShapeVector output_dims(shape.GetDims().begin(), shape.GetDims().end() - 1); - if (keepdims_) output_dims.push_back(1); - Tensor* output = context->Output(0, TensorShape(output_dims)); - const size_t partial_count = SafeInt(rows) * GetArgMaxLastAxisPartitions(width); - auto partial_values = GetScratchBuffer(partial_count, GetComputeStream(context)); - auto partial_indices = GetScratchBuffer(partial_count, GetComputeStream(context)); - using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; - return LaunchArgMaxLastAxis(Stream(context), reinterpret_cast(input->Data()), - input->GetElementType(), rows, width, - output->MutableData(), partial_values.get(), partial_indices.get()); -} - SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) : CudaKernel(info), head_dim_(narrow(info.GetAttrOrDefault("head_dim", 0))), diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu index 8221df1706c61..b9cf20a50001e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -17,71 +17,6 @@ namespace onnxruntime::contrib::cuda { namespace { constexpr int kThreads = 256; -constexpr int kArgmaxThreads = 256; - -__device__ __forceinline__ float ToFloat(float value) { return value; } -__device__ __forceinline__ float ToFloat(half value) { return __half2float(value); } -__device__ __forceinline__ float ToFloat(__nv_bfloat16 value) { return __bfloat162float(value); } - -template -__global__ void ArgmaxRowsKernel(const T* input, const int64_t* input_indices, - int64_t width, int64_t partitions, float* partial_values, - int64_t* output) { - __shared__ float values[kArgmaxThreads]; - __shared__ int64_t indices[kArgmaxThreads]; - const int tid = threadIdx.x; - const int64_t row = blockIdx.x / partitions; - const int64_t partition = blockIdx.x % partitions; - const int64_t end = partitions == 1 ? width : min(width, (partition + 1) * 4096); - float best_value = -CUDART_INF_F; - int64_t best_index = INT64_MAX; - for (int64_t col = partition * 4096 + tid; col < end; col += blockDim.x) { - const float value = ToFloat(input[row * width + col]); - const int64_t index = input_indices == nullptr ? col : input_indices[row * width + col]; - if (value > best_value || (value == best_value && index < best_index)) { - best_value = value; - best_index = index; - } - } - values[tid] = best_value; - indices[tid] = best_index; - __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - const float other_value = values[tid + stride]; - const int64_t other_index = indices[tid + stride]; - if (other_value > values[tid] || - (other_value == values[tid] && other_index < indices[tid])) { - values[tid] = other_value; - indices[tid] = other_index; - } - } - __syncthreads(); - } - if (tid == 0) { - output[blockIdx.x] = partial_values == nullptr && indices[0] == INT64_MAX ? 0 : indices[0]; - if (partial_values != nullptr) partial_values[blockIdx.x] = values[0]; - } -} - -template -Status LaunchTypedArgmax(cudaStream_t stream, const T* input, int64_t rows, - int64_t width, int64_t* output, - float* partial_values, int64_t* partial_indices) { - if (rows == 0) return Status::OK(); - const int64_t partitions = GetArgMaxLastAxisPartitions(width); - ORT_RETURN_IF_NOT(rows <= std::numeric_limits::max() / partitions, - "ArgMaxLastAxis has too many rows for a CUDA grid"); - ArgmaxRowsKernel<<(rows * partitions), kArgmaxThreads, 0, stream>>>( - input, nullptr, width, partitions, partitions == 1 ? nullptr : partial_values, - partitions == 1 ? output : partial_indices); - CUDA_RETURN_IF_ERROR(cudaGetLastError()); - if (partitions > 1) { - ArgmaxRowsKernel<<(rows), kArgmaxThreads, 0, stream>>>( - partial_values, partial_indices, partitions, 1, nullptr, output); - } - return CUDA_CALL(cudaGetLastError()); -} __device__ __forceinline__ float StableSigmoid(float value) { if (value >= 0.0f) { @@ -233,27 +168,6 @@ __global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention, } // namespace -Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_type, - int64_t rows, int64_t width, int64_t* output, - float* partial_values, int64_t* partial_indices) { - ORT_RETURN_IF_NOT(rows >= 0 && width > 0, - "ArgMaxLastAxis requires a non-negative row count and positive vocabulary dimension"); - if (rows == 0) return Status::OK(); - switch (input_type) { - case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: - return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); - case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: - return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); - case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: - return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); - default: - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "ArgMaxLastAxis input type is not supported"); - } -} - Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, const uint16_t* activation, const uint16_t* weight, uint16_t* prepared, float* row_gates, diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h index 4c57a089b886f..6d4a6a41b6041 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h @@ -9,12 +9,6 @@ namespace onnxruntime::contrib::cuda { -Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_type, - int64_t rows, int64_t width, int64_t* output, - float* partial_values, int64_t* partial_indices); -inline int64_t GetArgMaxLastAxisPartitions(int64_t width) { - return 1 + (width - 1) / 4096; -} Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, const uint16_t* projection, uint16_t* output, int64_t count, int head_dim, int heads, int q_stride, @@ -31,16 +25,6 @@ Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* prepared, float* row_gates, int rows, int width, int input_width); -template -class ArgMaxLastAxis final : public onnxruntime::cuda::CudaKernel { - public: - explicit ArgMaxLastAxis(const OpKernelInfo& info); - Status ComputeInternal(OpKernelContext* context) const override; - - private: - bool keepdims_; -}; - class SigmoidGateMul final : public onnxruntime::cuda::CudaKernel { public: explicit SigmoidGateMul(const OpKernelInfo& info); diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add.cc b/onnxruntime/contrib_ops/cuda/bert/gated_add.cc index 9480f08c43d30..e35c90078955a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add.cc @@ -3,6 +3,7 @@ #include "contrib_ops/cuda/bert/gated_add.h" #include "contrib_ops/cuda/bert/gated_add_impl.h" +#include "contrib_ops/cuda/bert/fused_tensor_ops.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" @@ -36,6 +37,7 @@ Status GatedAdd::ComputeInternal(OpKernelContext* context) const { const Tensor* x = context->Input(0); const Tensor* y = context->Input(1); const Tensor* gate = context->Input(2); + const Tensor* gate_weight = context->Input(3); const TensorShape& shape = x->Shape(); ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1, "X must have rank >= 1"); @@ -46,13 +48,32 @@ Status GatedAdd::ComputeInternal(OpKernelContext* context) const { const size_t last_axis = shape.NumDimensions() - 1; const int64_t hidden_size = shape[last_axis]; ORT_RETURN_IF_NOT(hidden_size > 0, "X last dimension must be positive"); - ORT_RETURN_IF_NOT(gate->Shape()[last_axis] == 1, "gate last dimension must be 1"); for (size_t axis = 0; axis < last_axis; ++axis) { ORT_RETURN_IF_NOT(gate->Shape()[axis] == shape[axis], "gate dimension ", axis, " must match X"); } Tensor* output = context->Output(0, shape); + if (gate_weight != nullptr) { + if constexpr (!std::is_same_v) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "GatedAdd gate_weight currently supports BFloat16 only"); + } else { + const int64_t input_width = gate->Shape()[last_axis]; + ORT_RETURN_IF_NOT(input_width > 0 && gate_weight->Shape().NumDimensions() == 2 && + gate_weight->Shape()[0] == input_width && gate_weight->Shape()[1] == 1, + "GatedAdd gate_weight must have shape [K, 1]"); + return LaunchProjectedSigmoidGateMul( + Stream(context), + reinterpret_cast(y->Data()), + reinterpret_cast(gate->Data()), + reinterpret_cast(gate_weight->Data()), + reinterpret_cast(x->Data()), + reinterpret_cast(output->MutableData()), + shape.Size() / hidden_size, input_width, hidden_size, sigmoid_in_fp32_); + } + } + + ORT_RETURN_IF_NOT(gate->Shape()[last_axis] == 1, "gate last dimension must be 1"); return LaunchGatedAddKernel( Stream(context), reinterpret_cast(output->MutableData()), diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add.h b/onnxruntime/contrib_ops/cuda/bert/gated_add.h index bd6f68e6d7147..2a48ab58769c3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add.h @@ -13,8 +13,16 @@ namespace cuda { template class GatedAdd final : public onnxruntime::cuda::CudaKernel { public: - explicit GatedAdd(const OpKernelInfo& info) : CudaKernel(info) {} + explicit GatedAdd(const OpKernelInfo& info) + : CudaKernel(info), + sigmoid_in_fp32_(info.GetAttrOrDefault("sigmoid_in_fp32", 1) == 1) { + const int64_t fp32 = info.GetAttrOrDefault("sigmoid_in_fp32", 1); + ORT_ENFORCE(fp32 == 0 || fp32 == 1, "GatedAdd sigmoid_in_fp32 must be zero or one"); + } Status ComputeInternal(OpKernelContext* context) const override; + + private: + bool sigmoid_in_fp32_{true}; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu index 21e1ff6983a38..9b061a4b555e7 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu @@ -61,6 +61,30 @@ __global__ void GatedAddKernel(T* output, } } +template +__global__ void GatedAddVec8Kernel(uint4* output, + const uint4* x, + const uint4* y, + const T* gate, + int vec_count, + onnxruntime::cuda::fast_divmod vec_hidden_size) { + const int v = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (v < vec_count) { + const T g = gate[vec_hidden_size.div(v)]; + const uint4 x_u4 = x[v]; + const uint4 y_u4 = y[v]; + uint4 out_u4; + const T* x_t = reinterpret_cast(&x_u4); + const T* y_t = reinterpret_cast(&y_u4); + T* out_t = reinterpret_cast(&out_u4); +#pragma unroll + for (int m = 0; m < 8; ++m) { + out_t[m] = RoundedMulAdd(x_t[m], y_t[m], g); + } + output[v] = out_u4; + } +} + } // namespace template @@ -76,6 +100,23 @@ Status LaunchGatedAddKernel(cudaStream_t stream, } constexpr int kThreads = 256; + if constexpr (sizeof(T) == 2) { + if ((hidden_size & 7) == 0 && (count >> 3) <= std::numeric_limits::max() && + (reinterpret_cast(output) & 15) == 0 && + (reinterpret_cast(x) & 15) == 0 && + (reinterpret_cast(y) & 15) == 0) { + const int vec_count = static_cast(count >> 3); + const int vec_blocks = (vec_count - 1) / kThreads + 1; + GatedAddVec8Kernel<<>>( + reinterpret_cast(output), + reinterpret_cast(x), + reinterpret_cast(y), + gate, + vec_count, + onnxruntime::cuda::fast_divmod(static_cast(hidden_size >> 3))); + return CUDA_CALL(cudaGetLastError()); + } + } const int64_t blocks = (count - 1) / kThreads + 1; ORT_RETURN_IF_NOT(blocks <= std::numeric_limits::max(), "GatedAdd launch requires too many blocks"); diff --git a/onnxruntime/contrib_ops/cuda/bert/scaled_silu_impl.cu b/onnxruntime/contrib_ops/cuda/bert/scaled_silu_impl.cu index f05bc281b975a..249c44d70ce11 100644 --- a/onnxruntime/contrib_ops/cuda/bert/scaled_silu_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/scaled_silu_impl.cu @@ -40,12 +40,43 @@ __global__ void ScaledSiLUKernel(const T* x, const void* scale, int scale_type, y[i] = from_float(to_float(z) * to_float(sigmoid)); } +template +__global__ void ScaledSiLUVec8Kernel(const uint4* x, const void* scale, int scale_type, + uint4* y, int vec_count, float alpha) { + const int v = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (v >= vec_count) return; + const float scale_value = scale == nullptr ? alpha : LoadScale(scale, scale_type); + const uint4 x_u4 = x[v]; + uint4 y_u4; + const T* x_t = reinterpret_cast(&x_u4); + T* y_t = reinterpret_cast(&y_u4); +#pragma unroll + for (int m = 0; m < 8; ++m) { + const T z = from_float(to_float(x_t[m]) * scale_value); + const T sigmoid = from_float(Sigmoid(to_float(z))); + y_t[m] = from_float(to_float(z) * to_float(sigmoid)); + } + y[v] = y_u4; +} + } // namespace template Status LaunchScaledSiLU(cudaStream_t stream, const T* x, const void* scale, int scale_type, T* y, int64_t count, float alpha) { if (count == 0) return Status::OK(); + if constexpr (sizeof(T) == 2) { + if ((count & 7) == 0 && (count >> 3) <= std::numeric_limits::max() && + (reinterpret_cast(x) & 15) == 0 && + (reinterpret_cast(y) & 15) == 0) { + const int vec_count = static_cast(count >> 3); + const int vec_blocks = (vec_count - 1) / kThreads + 1; + ScaledSiLUVec8Kernel<<>>( + reinterpret_cast(x), scale, scale_type, + reinterpret_cast(y), vec_count, alpha); + return CUDA_CALL(cudaGetLastError()); + } + } const int64_t block_count = (count - 1) / kThreads + 1; ORT_RETURN_IF_NOT(block_count <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); ScaledSiLUKernel<<(block_count), kThreads, 0, stream>>>(x, scale, scale_type, y, count, alpha); diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 6b4f1f0d21915..c75488f452caa 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -170,9 +170,6 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, LinearAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, GatedDeltaNet); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, GatedDeltaNet); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, GatedDeltaNet); -class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, ArgMaxLastAxis); -class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, ArgMaxLastAxis); -class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, ArgMaxLastAxis); class CUDA_MS_OP_CLASS_NAME(1, SigmoidGateMul); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttentionGate); @@ -541,9 +538,6 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 3f7415c32bd96..9739da6491154 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -5725,37 +5725,6 @@ constexpr const char* FusedTensorOps_ver1_doc = R"DOC( CUDA contrib operators for fused tensor operations used by transformer models. )DOC"; -ONNX_MS_OPERATOR_SET_SCHEMA( - ArgMaxLastAxis, 1, - OpSchema() - .SetDoc(FusedTensorOps_ver1_doc) - .Attr("keepdims", "Keep a singleton final axis when nonzero.", AttributeProto::INT, - static_cast(0)) - .Input(0, "logits", "Logits with vocabulary on the last axis.", "T") - .Output(0, "indices", - "First maximum index for each leading row. NaNs are ignored; " - "an all-NaN row returns index zero.", - "I") - .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Supported logit types.") - .TypeConstraint("I", {"tensor(int64)"}, "Argmax indices are int64.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - updateOutputElemType(ctx, 0, ONNX_NAMESPACE::TensorProto_DataType_INT64); - if (!hasInputShape(ctx, 0)) return; - const auto& input_shape = getInputShape(ctx, 0); - if (input_shape.dim_size() < 1) { - fail_shape_inference("ArgMaxLastAxis: logits must have rank at least one"); - } - ONNX_NAMESPACE::TensorShapeProto output_shape; - for (int i = 0; i + 1 < input_shape.dim_size(); ++i) { - *output_shape.add_dim() = input_shape.dim(i); - } - if (getAttribute(ctx, "keepdims", static_cast(0)) != 0) { - output_shape.add_dim()->set_dim_value(1); - } - updateOutputShape(ctx, 0, output_shape); - })); - ONNX_MS_OPERATOR_SET_SCHEMA( SigmoidGateMul, 1, OpSchema() @@ -6191,9 +6160,13 @@ ONNX_MS_OPERATOR_SET_SCHEMA( GatedAdd, 1, OpSchema() .SetDoc(GatedAdd_ver1_doc) + .Attr("sigmoid_in_fp32", + "When gate_weight is provided, compute sigmoid in FP32 (1, default) or input type (0).", + AttributeProto::INT, static_cast(1)) .Input(0, "X", "Unscaled input with shape (..., C).", "T") .Input(1, "Y", "Input scaled by gate, with the same shape as X.", "T") - .Input(2, "gate", "Per-row gate with shape (..., 1).", "T") + .Input(2, "gate", "Per-row gate with shape (..., 1), or activation (..., K) when gate_weight is provided.", "T") + .Input(3, "gate_weight", "Optional projection weight [K, 1] applied to gate before sigmoid.", "T", OpSchema::Optional) .Output(0, "output", "Gated sum with the same shape as X.", "T") .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 7696c0901f3a3..53842f1666c0f 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -97,7 +97,6 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPostMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionLowRankMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNet); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNetStateReplay); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ArgMaxLastAxis); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SigmoidGateMul); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ConvStateReplay); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedRMSNorm); @@ -235,7 +234,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/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 67bf96d8f3199..1c1ca9088b7a5 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -1392,9 +1392,11 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kO class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, float, ArgMax); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, double, ArgMax); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, MLFloat16, ArgMax); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ArgMax); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, float, ArgMin); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, double, ArgMin); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, MLFloat16, ArgMin); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ArgMin); // OpSet 14 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, CumSum); @@ -2722,9 +2724,11 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, // OpSet 14 BuildKernelCreateInfo, diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index 60ffd772c97f1..5fc35f5e68858 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -777,6 +777,12 @@ struct ArgReductionIdentity { __device__ __forceinline__ static half Highest() { return __ushort_as_half(static_cast(0x7C00U)); } }; +template <> +struct ArgReductionIdentity { + __device__ __forceinline__ static BFloat16 Lowest() { return BFloat16::FromBits(static_cast(0xFF80U)); } + __device__ __forceinline__ static BFloat16 Highest() { return BFloat16::FromBits(static_cast(0x7F80U)); } +}; + template __device__ __forceinline__ T arg_reduction_identity() { return IsArgMax ? ArgReductionIdentity::Lowest() : ArgReductionIdentity::Highest(); @@ -1163,6 +1169,7 @@ Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* out int64_t* output, int m, int n, void* buffer, \ size_t buffer_size) INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(half); +INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(BFloat16); INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(float); INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(double); #undef INSTANTIATE_ARG_MIN_MAX_LAST_AXIS diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index 3d919ac04ecbf..a110cb1e55ce2 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -97,8 +97,9 @@ Status ReduceKernel::ReduceKernelShared( int64_t input_count = input_shape.Size(); IAllocatorUniquePtr temp_X; - if (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES && std::is_same::value) { - // ArgMax/ArgMin with FP16 are not supported by cudnn, so convert input to fp32 then call cudnn + if (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES && + (std::is_same::value || std::is_same::value)) { + // ArgMax/ArgMin with FP16/BF16 are not supported by cudnn, so convert input to fp32 then call cudnn temp_X = GetScratchBuffer(input_count, stream); cudnn_type_X = CUDNN_DATA_FLOAT; Impl_Cast(cuda_stream, reinterpret_cast(X), temp_X.get(), input_shape.Size()); @@ -114,7 +115,7 @@ Status ReduceKernel::ReduceKernelShared( } CudnnReduceDescriptor reduce_desc; - if constexpr (std::is_same::value) + if constexpr (std::is_same::value || std::is_same::value) ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, CudnnTensor::GetDataType(), ReduceTensorIndices)); else ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, cudnn_type_X, ReduceTensorIndices)); @@ -527,9 +528,10 @@ Status ReduceComputeCore(const AllocatorPtr& gpu_allocator, const CudaKernel* ke cudnnDataType_t cudnn_type_X = CUDNN_DATA_FLOAT; // Reducesum with BFP16 is not supported by cudnn, so convert input to fp32 then call cudnn - if ((ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES && std::is_same::value) || + if ((ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES && + (std::is_same::value || std::is_same::value)) || (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_NO_INDICES && std::is_same::value)) { - // ArgMax/ArgMin with FP16 are not supported by cudnn, so convert input to fp32 then call cudnn + // ArgMax/ArgMin with FP16/BF16 are not supported by cudnn, so convert input to fp32 then call cudnn temp_X = AllocateScratchBuffer(gpu_allocator, kernel, input_count, compute_stream); Impl_Cast(stream, reinterpret_cast(input.Data()), temp_X.get(), input_shape.Size()); } else { @@ -1031,10 +1033,12 @@ template std::unique_ptr ReduceCompute("axis", -1); + test.AddAttribute("keepdims", 0); test.AddInput("logits", {2, 4}, {1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f}); test.AddOutput("indices", {2}, {1, 0}); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); @@ -49,7 +51,8 @@ TEST(FusedTensorOpsTest, ArgMaxLastAxisWithKeepdims) { auto providers = CudaProvider(); if (providers.empty()) GTEST_SKIP(); - OpTester test("ArgMaxLastAxis", 1, kMSDomain); + OpTester test("ArgMax", 13); + test.AddAttribute("axis", -1); test.AddAttribute("keepdims", 1); test.AddInput("logits", {2, 4}, ToTensorType({1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f})); @@ -65,7 +68,8 @@ TEST(FusedTensorOpsTest, ArgMaxLastAxisPartitionedVocabulary) { std::vector values(3 * width, -std::numeric_limits::infinity()); values[4095] = values[4096] = values[width - 1] = 5.0f; values[2 * width - 1] = 7.0f; - OpTester test("ArgMaxLastAxis", 1, kMSDomain); + OpTester test("ArgMax", 13); + test.AddAttribute("axis", -1); test.AddAttribute("keepdims", keepdims ? 1 : 0); test.AddInput("logits", {3, width}, ToBFloat16(values)); test.AddOutput("indices", keepdims ? std::vector{3, 1} : std::vector{3}, @@ -81,12 +85,16 @@ void RunArgMaxNaNRows() { if (providers.empty()) GTEST_SKIP(); const float nan = std::numeric_limits::quiet_NaN(); std::vector values(4 * width, nan); + values[width] = -std::numeric_limits::infinity(); values[2 * width - 1] = -std::numeric_limits::infinity(); values[2 * width] = values[3 * width - 1] = 5.0f; + values[3 * width] = -1.0f; values[4 * width - 1] = 7.0f; - OpTester test("ArgMaxLastAxis", 1, kMSDomain); + OpTester test("ArgMax", 13); + test.AddAttribute("axis", -1); + test.AddAttribute("keepdims", 0); test.AddInput("logits", {4, width}, ToTensorType(values)); - test.AddOutput("indices", {4}, {0, width - 1, 0, width - 1}); + test.AddOutput("indices", {4}, {0, 0, 0, width - 1}); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } } From 91b3096fa8367ff4d3b78e9c5f29447c4d01c08a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 25 Oct 2026 02:58:23 +0000 Subject: [PATCH 14/19] fix(cuda): remove leftover ArgMaxLastAxis template instantiations Remove the three leftover explicit template instantiations of ArgMaxLastAxis in fused_tensor_ops.cc after ArgMaxLastAxis was replaced with BFloat16 support on ai.onnx::ArgMax and ai.onnx::ArgMin. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc | 4 ---- 1 file changed, 4 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc index ff1544dcf0c03..28d5447e4a029 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -127,8 +127,4 @@ Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { dense_prefill ? false : sigmoid_in_fp32_, silu_gate_); } -template class ArgMaxLastAxis; -template class ArgMaxLastAxis; -template class ArgMaxLastAxis; - } // namespace onnxruntime::contrib::cuda From 10fbe33795d0ba07231aefa5ede96534259266be Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 25 Oct 2026 05:01:19 +0000 Subject: [PATCH 15/19] Remove unused SigmoidGateMul contrib op and move projected gate kernel into GatedAdd - Move LaunchProjectedSigmoidGateMul directly into gated_add_impl.{h,cu} for GatedAdd's optional gate_weight input. - Remove unused com.microsoft::SigmoidGateMul schema, kernel registration, fused_tensor_ops.{h,cc,cu}, and docs. - Add BFloat16 projected-gate GatedAdd and ArgMax unit tests to gated_add_op_test.cc. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 62 ---- docs/OperatorKernels.md | 1 - docs/contrib_ops/cuda/sigmoid_gate_mul.md | 25 -- .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 130 -------- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 240 -------------- .../contrib_ops/cuda/bert/fused_tensor_ops.h | 44 --- .../contrib_ops/cuda/bert/gated_add.cc | 1 - .../contrib_ops/cuda/bert/gated_add_impl.cu | 109 +++++++ .../contrib_ops/cuda/bert/gated_add_impl.h | 11 + .../contrib_ops/cuda/bert/scaled_silu_impl.cu | 31 -- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 2 - .../core/graph/contrib_ops/bert_defs.cc | 42 --- onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 - .../test/contrib_ops/fused_tensor_ops_test.cc | 302 ------------------ .../test/contrib_ops/gated_add_op_test.cc | 76 +++++ 15 files changed, 196 insertions(+), 882 deletions(-) delete mode 100644 docs/contrib_ops/cuda/sigmoid_gate_mul.md delete mode 100644 onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc delete mode 100644 onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu delete mode 100644 onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h delete mode 100644 onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index bdcc34f44971e..c33e0c4eefc96 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -123,7 +123,6 @@ Do not modify directly.* * com.microsoft.SampleOp * com.microsoft.Sampling * com.microsoft.ScaledSiLU - * com.microsoft.SigmoidGateMul * com.microsoft.SkipGroupNorm * com.microsoft.SkipLayerNormalization * com.microsoft.SkipSimplifiedLayerNormalization @@ -7569,67 +7568,6 @@ This version of the operator has been available since version 1 of the 'com.micr -### **com.microsoft.SigmoidGateMul** - - CUDA contrib operators for fused tensor operations used by transformer models. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
broadcast_gate : int
-
When one, projection has the attention shape with final dimension one. Packed layout attributes must be zero. CUDA only.
-
head_dim : int
-
Dimension of each attention head.
-
num_heads : int
-
Number of attention heads.
-
q_offset : int
-
Offset of the query projection within each token.
-
q_stride : int
-
Per-head stride in the packed projection.
-
sigmoid_in_fp32 : int
-
One computes sigmoid in FP32 before BF16 conversion. Zero preserves the input-type intermediate arithmetic of CUDA Sigmoid, including subtraction on negative inputs.
-
silu_gate : int
-
When one, projection and attention have identical shapes. Multiply the BF16-rounded sigmoid by projection and round to BF16 before multiplying attention. Cannot be combined with broadcast_gate or packed layout attributes.
-
tok_stride : int
-
Per-token stride in the packed projection.
-
- -#### Inputs (2 - 5) - -
-
attention : T
-
BF16 attention values.
-
projection : T
-
BF16 packed projection containing gate values.
-
residual (optional) : T
-
Optional BF16 residual with the attention shape.
-
prefill_numerics (optional) : TI
-
Optional CPU int32 scalar 0/1.
-
gate_weight (optional) : T
-
Optional BF16 projected-gate weight [input_width,1], requires broadcast_gate.
-
- -#### Outputs - -
-
output : T
-
Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 before optionally adding residual and rounding again.
-
- -#### Type Constraints - -
-
T : tensor(bfloat16)
-
Inputs and output use BF16.
-
TI : tensor(int32)
-
CPU numerical control.
-
- - ### **com.microsoft.SkipGroupNorm** This operator element-wise adds x, skip and bias, then apply group normalization and optional activation. diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index a6628392f0837..fc416deb384c4 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1171,7 +1171,6 @@ The **OpSet Version** column uses the following notation: |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)| -|SigmoidGateMul|*in* attention:**T**
*in* projection:**T**
*in* residual:**T**
*in* prefill_numerics:**TI**
*in* gate_weight:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16)
**TI** = tensor(int32)| |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)| |SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |SkipSimplifiedLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| diff --git a/docs/contrib_ops/cuda/sigmoid_gate_mul.md b/docs/contrib_ops/cuda/sigmoid_gate_mul.md deleted file mode 100644 index 7e57d370f77ec..0000000000000 --- a/docs/contrib_ops/cuda/sigmoid_gate_mul.md +++ /dev/null @@ -1,25 +0,0 @@ -# Projected sigmoid gates and shared-expert preparation - -The CUDA BF16 `SigmoidGateMul` operator retains its existing packed, broadcast -and SiLU modes. Two appended optional inputs support shared-expert fusion: - -- Input 3, `prefill_numerics`, is a CPU INT32 tensor with one value, 0 or 1. - Mode 1 selects BF16 sigmoid intermediates; omitting it preserves the existing - `sigmoid_in_fp32` attribute. -- Input 4, `gate_weight`, is BF16 `[input_width, 1]`. It requires - `broadcast_gate=1` and `silu_gate=0`. In this mode input 1 is the projection - activation with the same leading dimensions as input 0, rather than an - already projected scalar gate. Its final dimension is `input_width`. - -The projected-gate implementation computes the dot product in FP32 using a -256-thread CUB reduction, rounds to BF16, applies the selected sigmoid -arithmetic, and preserves the existing BF16 multiply/residual rounding. -The residual, when supplied, must match input 0. - -`LaunchPrepareSharedExpert` shares the same projected-gate reduction and BF16 -sigmoid arithmetic. One launch prepares the SiLU/up activation and the scalar -row gates for `MatMulBlockQuantizedFp8Weight(shared_expert=1)` in companion -PR #33473. No behavior changes occur without the new inputs or opt-in mode. - -The shared-expert native tests compare these paths bit-exactly with the -unfused CUDA operator sequence, including prefill and fallback dispatch. diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc deleted file mode 100644 index 28d5447e4a029..0000000000000 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ /dev/null @@ -1,130 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "contrib_ops/cuda/bert/fused_tensor_ops.h" - -#include -#include -#include - -#include "core/common/narrow.h" -#include "core/common/safeint.h" -#include "core/providers/cuda/cuda_common.h" -#include "core/providers/cuda/cuda_type_conversion.h" - -namespace onnxruntime::contrib::cuda { - -ONNX_OPERATOR_KERNEL_EX( - SigmoidGateMul, kMSDomain, 1, kCudaExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()) - .TypeConstraint("TI", DataTypeImpl::GetTensorType()) - .InputMemoryType(OrtMemTypeCPUInput, 3), - SigmoidGateMul); - -SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) - : CudaKernel(info), - head_dim_(narrow(info.GetAttrOrDefault("head_dim", 0))), - heads_(narrow(info.GetAttrOrDefault("num_heads", 0))), - q_stride_(narrow(info.GetAttrOrDefault("q_stride", 0))), - tok_stride_(info.GetAttrOrDefault("tok_stride", 0)), - q_offset_(info.GetAttrOrDefault("q_offset", 0)), - broadcast_gate_(info.GetAttrOrDefault("broadcast_gate", 0) == 1), - sigmoid_in_fp32_(info.GetAttrOrDefault("sigmoid_in_fp32", 1) == 1), - silu_gate_(info.GetAttrOrDefault("silu_gate", 0) == 1) { - const int64_t broadcast = info.GetAttrOrDefault("broadcast_gate", 0); - ORT_ENFORCE(broadcast == 0 || broadcast == 1, "SigmoidGateMul broadcast_gate must be zero or one"); - const int64_t fp32 = info.GetAttrOrDefault("sigmoid_in_fp32", 1); - ORT_ENFORCE(fp32 == 0 || fp32 == 1, "SigmoidGateMul sigmoid_in_fp32 must be zero or one"); - const int64_t silu = info.GetAttrOrDefault("silu_gate", 0); - ORT_ENFORCE(silu == 0 || silu == 1, "SigmoidGateMul silu_gate must be zero or one"); - ORT_ENFORCE(!(broadcast_gate_ && silu_gate_), "SigmoidGateMul broadcast_gate and silu_gate are exclusive"); - ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && q_stride_ >= 2LL * head_dim_ && tok_stride_ > 0 && q_offset_ >= 0), - "SigmoidGateMul received invalid layout attributes"); - ORT_ENFORCE(!(broadcast_gate_ || silu_gate_) || (head_dim_ == 0 && heads_ == 0 && q_stride_ == 0 && - tok_stride_ == 0 && q_offset_ == 0), - "SigmoidGateMul broadcast_gate/silu_gate cannot be combined with packed layout attributes"); -} - -Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { - const Tensor* numerics = context->Input(3); - 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; - const Tensor* attention = context->Input(0); - const Tensor* projection = context->Input(1); - const Tensor* residual = context->Input(2); - const Tensor* gate_weight = context->Input(4); - const int64_t count = attention->Shape().Size(); - ORT_RETURN_IF_NOT(residual == nullptr || residual->Shape() == attention->Shape(), - "SigmoidGateMul residual must have the same shape as attention"); - if (gate_weight != nullptr) { - ORT_RETURN_IF_NOT(broadcast_gate_ && !silu_gate_, - "SigmoidGateMul gate_weight requires broadcast_gate without silu_gate"); - const auto& shape = attention->Shape(); - const auto& activation_shape = projection->Shape(); - ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1 && shape[shape.NumDimensions() - 1] > 0 && - activation_shape.NumDimensions() == shape.NumDimensions() && - activation_shape[activation_shape.NumDimensions() - 1] > 0, - "SigmoidGateMul projected gate requires matching ranks and positive final dimensions"); - for (size_t axis = 0; axis + 1 < shape.NumDimensions(); ++axis) { - ORT_RETURN_IF_NOT(shape[axis] == activation_shape[axis], - "SigmoidGateMul projected gate leading dimensions must match attention"); - } - const int64_t input_width = activation_shape[activation_shape.NumDimensions() - 1]; - ORT_RETURN_IF_NOT(gate_weight->Shape().NumDimensions() == 2 && - gate_weight->Shape()[0] == input_width && gate_weight->Shape()[1] == 1, - "SigmoidGateMul gate_weight must have shape [input_width, 1]"); - Tensor* output = context->Output(0, shape); - return LaunchProjectedSigmoidGateMul( - Stream(context), reinterpret_cast(attention->Data()), - reinterpret_cast(projection->Data()), - reinterpret_cast(gate_weight->Data()), - residual == nullptr ? nullptr : reinterpret_cast(residual->Data()), - reinterpret_cast(output->MutableData()), - count / shape[shape.NumDimensions() - 1], input_width, shape[shape.NumDimensions() - 1], - dense_prefill ? false : sigmoid_in_fp32_); - } - int64_t gate_row_width = 0; - if (silu_gate_) { - ORT_RETURN_IF_NOT(projection->Shape() == attention->Shape(), - "SigmoidGateMul silu gate must have the same shape as attention"); - } else if (broadcast_gate_) { - TensorShapeVector expected(attention->Shape().GetDims().begin(), attention->Shape().GetDims().end()); - ORT_RETURN_IF_NOT(!expected.empty() && expected.back() > 0, - "SigmoidGateMul broadcast_gate requires a positive final attention dimension"); - gate_row_width = expected.back(); - expected.back() = 1; - ORT_RETURN_IF_NOT(projection->Shape() == TensorShape(expected), - "SigmoidGateMul broadcast gate must match attention with a singleton final dimension"); - } else { - const int64_t row_width = static_cast(heads_) * head_dim_; - ORT_RETURN_IF_NOT(count % row_width == 0, - "SigmoidGateMul attention size must be divisible by heads * head_dim"); - const int64_t tokens = count / row_width; - if (tokens > 0) { - int64_t required_extent; - int64_t head_extent; - ORT_RETURN_IF_NOT(SafeMultiply(tokens - 1, tok_stride_, required_extent) && - SafeMultiply(static_cast(heads_ - 1), static_cast(q_stride_), head_extent) && - SafeAdd(head_extent, 2LL * head_dim_, head_extent) && - SafeAdd(required_extent, q_offset_, required_extent) && - SafeAdd(required_extent, head_extent, required_extent), - "SigmoidGateMul projection extent overflows int64"); - ORT_RETURN_IF_NOT(projection->Shape().Size() >= required_extent, - "SigmoidGateMul projection is too small for the configured layout"); - } - } - Tensor* output = context->Output(0, attention->Shape()); - return LaunchSigmoidGateMul( - Stream(context), reinterpret_cast(attention->Data()), - reinterpret_cast(projection->Data()), - reinterpret_cast(output->MutableData()), count, head_dim_, heads_, - q_stride_, tok_stride_, q_offset_, gate_row_width, - residual == nullptr ? nullptr : reinterpret_cast(residual->Data()), - dense_prefill ? false : sigmoid_in_fp32_, silu_gate_); -} - -} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu deleted file mode 100644 index b9cf20a50001e..0000000000000 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ /dev/null @@ -1,240 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "contrib_ops/cuda/bert/fused_tensor_ops.h" - -#include -#include -#include -#include -#include -#include - -#include "core/providers/cuda/cuda_common.h" -#include "core/providers/cuda/cu_inc/common.cuh" - -namespace onnxruntime::contrib::cuda { -namespace { - -constexpr int kThreads = 256; - -__device__ __forceinline__ float StableSigmoid(float value) { - if (value >= 0.0f) { - return 1.0f / (1.0f + expf(-value)); - } - const float exp_value = expf(value); - return exp_value / (1.0f + exp_value); -} - -__device__ __forceinline__ float RoundedBf16Sigmoid(float value, bool sigmoid_in_fp32) { - if (sigmoid_in_fp32) { - return __bfloat162float(__float2bfloat16_rn(StableSigmoid(value))); - } - using namespace onnxruntime::cuda; - const BFloat16 a(value); - return static_cast(a > BFloat16(0.0f) - ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) - : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a))); -} - -__device__ float ProjectedRowGate(const __nv_bfloat16* activation, const __nv_bfloat16* weight, - int64_t row, int64_t input_width, bool sigmoid_in_fp32) { - using Reduce = cub::BlockReduce; - __shared__ typename Reduce::TempStorage reduction; - __shared__ float gate; - float dot = 0.0f; - for (int64_t k = threadIdx.x; k < input_width; k += kThreads) { - dot = fmaf(__bfloat162float(activation[row * input_width + k]), - __bfloat162float(weight[k]), dot); - } - const float total = Reduce(reduction).Sum(dot); - if (threadIdx.x == 0) { - const float value = __bfloat162float(__float2bfloat16_rn(total)); - gate = RoundedBf16Sigmoid(value, sigmoid_in_fp32); - } - __syncthreads(); - return gate; -} - -__global__ void PrepareSharedExpertKernel( - const __nv_bfloat16* up, const __nv_bfloat16* silu_gate, const __nv_bfloat16* activation, - const __nv_bfloat16* weight, __nv_bfloat16* prepared, float* row_gates, - int width, int input_width) { - const int row = blockIdx.x; - const float gate = ProjectedRowGate(activation, weight, row, input_width, false); - if (threadIdx.x == 0) row_gates[row] = gate; - for (int column = threadIdx.x; column < width; column += blockDim.x) { - const int64_t index = static_cast(row) * width + column; - const float value = __bfloat162float(silu_gate[index]); - const auto silu = __float2bfloat16_rn(value * RoundedBf16Sigmoid(value, false)); - prepared[index] = __float2bfloat16_rn(__bfloat162float(up[index]) * __bfloat162float(silu)); - } -} - -__global__ void ProjectedSigmoidGateMulKernel( - const __nv_bfloat16* attention, const __nv_bfloat16* activation, - const __nv_bfloat16* weight, const __nv_bfloat16* residual, __nv_bfloat16* output, - int64_t input_width, int64_t output_width, bool sigmoid_in_fp32) { - const int64_t row = blockIdx.x; - const float gate = ProjectedRowGate(activation, weight, row, input_width, sigmoid_in_fp32); - const int64_t row_base = row * output_width; - if ((output_width & 7) == 0 && - (reinterpret_cast(attention + row_base) & 15) == 0 && - (reinterpret_cast(output + row_base) & 15) == 0 && - (residual == nullptr || (reinterpret_cast(residual + row_base) & 15) == 0)) { - const int64_t vec_width = output_width >> 3; - const uint4* attn_vec = reinterpret_cast(attention + row_base); - const uint4* res_vec = residual != nullptr ? reinterpret_cast(residual + row_base) : nullptr; - uint4* out_vec = reinterpret_cast(output + row_base); - for (int64_t v = threadIdx.x; v < vec_width; v += kThreads) { - const uint4 a_u4 = attn_vec[v]; - const __nv_bfloat16* a_bf = reinterpret_cast(&a_u4); - uint4 o_u4; - __nv_bfloat16* o_bf = reinterpret_cast<__nv_bfloat16*>(&o_u4); - if (res_vec != nullptr) { - const uint4 r_u4 = res_vec[v]; - const __nv_bfloat16* r_bf = reinterpret_cast(&r_u4); -#pragma unroll - for (int m = 0; m < 8; ++m) { - const __nv_bfloat16 scaled = __float2bfloat16_rn(__bfloat162float(a_bf[m]) * gate); - o_bf[m] = __float2bfloat16_rn(__bfloat162float(scaled) + __bfloat162float(r_bf[m])); - } - } else { -#pragma unroll - for (int m = 0; m < 8; ++m) { - o_bf[m] = __float2bfloat16_rn(__bfloat162float(a_bf[m]) * gate); - } - } - out_vec[v] = o_u4; - } - return; - } - for (int64_t column = threadIdx.x; column < output_width; column += kThreads) { - const int64_t index = row_base + column; - __nv_bfloat16 value = __float2bfloat16_rn(__bfloat162float(attention[index]) * gate); - if (residual != nullptr) { - value = __float2bfloat16_rn(__bfloat162float(value) + __bfloat162float(residual[index])); - } - output[index] = value; - } -} - -template -__global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention, - const __nv_bfloat16* projection, - __nv_bfloat16* output, int64_t count, - int head_dim, int heads, int q_stride, - int64_t token_stride, int64_t q_offset, - int64_t gate_row_width, const __nv_bfloat16* residual, - bool sigmoid_in_fp32) { - const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= count) return; - int64_t gate_index; - if constexpr (SiluGate) { - gate_index = i; - } else if constexpr (BroadcastGate) { - gate_index = i / gate_row_width; - } else { - const int64_t row = i / head_dim; - const int64_t token = row / heads; - const int head = static_cast(row % heads); - const int dim = static_cast(i % head_dim); - gate_index = token * token_stride + q_offset + - static_cast(head) * q_stride + head_dim + dim; - } - const float gate = __bfloat162float(projection[gate_index]); - __nv_bfloat16 sigmoid_bf16; - if (sigmoid_in_fp32) { - sigmoid_bf16 = __float2bfloat16_rn(StableSigmoid(gate)); - } else { - using namespace onnxruntime::cuda; - const BFloat16 a(gate); - // Match CUDA Sigmoid's input-type intermediates, including its negative-input subtraction. - const BFloat16 value = a > BFloat16(0.0f) - ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) - : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a)); - sigmoid_bf16 = __float2bfloat16_rn(static_cast(value)); - } - if constexpr (SiluGate) { - sigmoid_bf16 = __float2bfloat16_rn(gate * __bfloat162float(sigmoid_bf16)); - } - __nv_bfloat16 value = __float2bfloat16_rn( - __bfloat162float(attention[i]) * __bfloat162float(sigmoid_bf16)); - if constexpr (HasResidual) { - value = __float2bfloat16_rn(__bfloat162float(value) + __bfloat162float(residual[i])); - } - output[i] = value; -} - -} // namespace - -Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, - const uint16_t* activation, const uint16_t* weight, - uint16_t* prepared, float* row_gates, - int rows, int width, int input_width) { - ORT_RETURN_IF_NOT(rows >= 0 && width > 0 && input_width > 0, - "Shared expert preparation requires nonnegative rows and positive widths"); - if (rows == 0) return Status::OK(); - PrepareSharedExpertKernel<<>>( - reinterpret_cast(up), reinterpret_cast(silu_gate), - reinterpret_cast(activation), reinterpret_cast(weight), - reinterpret_cast<__nv_bfloat16*>(prepared), row_gates, width, input_width); - return CUDA_CALL(cudaGetLastError()); -} - -Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, - const uint16_t* activation, const uint16_t* weight, - const uint16_t* residual, uint16_t* output, - int64_t rows, int64_t input_width, int64_t output_width, - bool sigmoid_in_fp32) { - ORT_RETURN_IF_NOT(rows >= 0 && rows <= std::numeric_limits::max() && - input_width > 0 && output_width > 0, - "SigmoidGateMul projected gate dimensions exceed CUDA limits"); - if (rows == 0) return Status::OK(); - ProjectedSigmoidGateMulKernel<<(rows), kThreads, 0, stream>>>( - reinterpret_cast(attention), - reinterpret_cast(activation), - reinterpret_cast(weight), - reinterpret_cast(residual), reinterpret_cast<__nv_bfloat16*>(output), - input_width, output_width, sigmoid_in_fp32); - return CUDA_CALL(cudaGetLastError()); -} - -Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, - const uint16_t* projection, uint16_t* output, int64_t count, - int head_dim, int heads, int q_stride, - int64_t token_stride, int64_t q_offset, - int64_t gate_row_width, const uint16_t* residual, bool sigmoid_in_fp32, - bool silu_gate) { - if (count == 0) return Status::OK(); - const int64_t block_count = (count - 1) / kThreads + 1; - ORT_RETURN_IF_NOT(block_count <= std::numeric_limits::max(), - "SigmoidGateMul requires too many CUDA blocks"); -#define LAUNCH_GATE(BROADCAST, RESIDUAL, SILU) \ - SigmoidGateMulKernel<<(block_count), kThreads, 0, stream>>>( \ - reinterpret_cast(attention), \ - reinterpret_cast(projection), \ - reinterpret_cast<__nv_bfloat16*>(output), count, head_dim, heads, q_stride, \ - token_stride, q_offset, gate_row_width, reinterpret_cast(residual), sigmoid_in_fp32) - if (silu_gate) { - if (residual != nullptr) { - LAUNCH_GATE(false, true, true); - } else { - LAUNCH_GATE(false, false, true); - } - } else if (gate_row_width > 0) { - if (residual != nullptr) { - LAUNCH_GATE(true, true, false); - } else { - LAUNCH_GATE(true, false, false); - } - } else if (residual != nullptr) { - LAUNCH_GATE(false, true, false); - } else { - LAUNCH_GATE(false, false, false); - } -#undef LAUNCH_GATE - return CUDA_CALL(cudaGetLastError()); -} - -} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h deleted file mode 100644 index 6d4a6a41b6041..0000000000000 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h +++ /dev/null @@ -1,44 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include - -#include "core/providers/cuda/cuda_kernel.h" - -namespace onnxruntime::contrib::cuda { - -Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, - const uint16_t* projection, uint16_t* output, int64_t count, - int head_dim, int heads, int q_stride, - int64_t token_stride, int64_t q_offset, - int64_t gate_row_width, const uint16_t* residual, bool sigmoid_in_fp32, - bool silu_gate); -Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, - const uint16_t* activation, const uint16_t* weight, - const uint16_t* residual, uint16_t* output, - int64_t rows, int64_t input_width, int64_t output_width, - bool sigmoid_in_fp32); -Status LaunchPrepareSharedExpert(cudaStream_t stream, const uint16_t* up, const uint16_t* silu_gate, - const uint16_t* activation, const uint16_t* weight, - uint16_t* prepared, float* row_gates, - int rows, int width, int input_width); - -class SigmoidGateMul final : public onnxruntime::cuda::CudaKernel { - public: - explicit SigmoidGateMul(const OpKernelInfo& info); - Status ComputeInternal(OpKernelContext* context) const override; - - private: - int head_dim_; - int heads_; - int q_stride_; - int64_t tok_stride_; - int64_t q_offset_; - bool broadcast_gate_; - bool sigmoid_in_fp32_; - bool silu_gate_; -}; - -} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add.cc b/onnxruntime/contrib_ops/cuda/bert/gated_add.cc index e35c90078955a..033a42c6f0673 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add.cc @@ -3,7 +3,6 @@ #include "contrib_ops/cuda/bert/gated_add.h" #include "contrib_ops/cuda/bert/gated_add_impl.h" -#include "contrib_ops/cuda/bert/fused_tensor_ops.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu index 9b061a4b555e7..adac73eb5789b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu @@ -4,11 +4,14 @@ #include #include #include +#include +#include #include #include "contrib_ops/cuda/bert/gated_add_impl.h" #include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/shared_inc/fast_divmod.h" namespace onnxruntime { @@ -17,6 +20,94 @@ namespace cuda { namespace { +constexpr int kProjectedGateThreads = 256; + +__device__ __forceinline__ float StableSigmoid(float value) { + if (value >= 0.0f) { + return 1.0f / (1.0f + expf(-value)); + } + const float exp_value = expf(value); + return exp_value / (1.0f + exp_value); +} + +__device__ __forceinline__ float RoundedBf16Sigmoid(float value, bool sigmoid_in_fp32) { + if (sigmoid_in_fp32) { + return __bfloat162float(__float2bfloat16_rn(StableSigmoid(value))); + } + using namespace onnxruntime::cuda; + const BFloat16 a(value); + return static_cast(a > BFloat16(0.0f) + ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) + : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a))); +} + +__device__ float ProjectedRowGate(const __nv_bfloat16* activation, const __nv_bfloat16* weight, + int64_t row, int64_t input_width, bool sigmoid_in_fp32) { + using Reduce = cub::BlockReduce; + __shared__ typename Reduce::TempStorage reduction; + __shared__ float gate; + float dot = 0.0f; + for (int64_t k = threadIdx.x; k < input_width; k += kProjectedGateThreads) { + dot = fmaf(__bfloat162float(activation[row * input_width + k]), + __bfloat162float(weight[k]), dot); + } + const float total = Reduce(reduction).Sum(dot); + if (threadIdx.x == 0) { + const float value = __bfloat162float(__float2bfloat16_rn(total)); + gate = RoundedBf16Sigmoid(value, sigmoid_in_fp32); + } + __syncthreads(); + return gate; +} + +__global__ void ProjectedSigmoidGateMulKernel( + const __nv_bfloat16* attention, const __nv_bfloat16* activation, + const __nv_bfloat16* weight, const __nv_bfloat16* residual, __nv_bfloat16* output, + int64_t input_width, int64_t output_width, bool sigmoid_in_fp32) { + const int64_t row = blockIdx.x; + const float gate = ProjectedRowGate(activation, weight, row, input_width, sigmoid_in_fp32); + const int64_t row_base = row * output_width; + if ((output_width & 7) == 0 && + (reinterpret_cast(attention + row_base) & 15) == 0 && + (reinterpret_cast(output + row_base) & 15) == 0 && + (residual == nullptr || (reinterpret_cast(residual + row_base) & 15) == 0)) { + const int64_t vec_width = output_width >> 3; + const uint4* attn_vec = reinterpret_cast(attention + row_base); + const uint4* res_vec = residual != nullptr ? reinterpret_cast(residual + row_base) : nullptr; + uint4* out_vec = reinterpret_cast(output + row_base); + for (int64_t v = threadIdx.x; v < vec_width; v += kProjectedGateThreads) { + const uint4 a_u4 = attn_vec[v]; + const __nv_bfloat16* a_bf = reinterpret_cast(&a_u4); + uint4 o_u4; + __nv_bfloat16* o_bf = reinterpret_cast<__nv_bfloat16*>(&o_u4); + if (res_vec != nullptr) { + const uint4 r_u4 = res_vec[v]; + const __nv_bfloat16* r_bf = reinterpret_cast(&r_u4); +#pragma unroll + for (int m = 0; m < 8; ++m) { + const __nv_bfloat16 scaled = __float2bfloat16_rn(__bfloat162float(a_bf[m]) * gate); + o_bf[m] = __float2bfloat16_rn(__bfloat162float(scaled) + __bfloat162float(r_bf[m])); + } + } else { +#pragma unroll + for (int m = 0; m < 8; ++m) { + o_bf[m] = __float2bfloat16_rn(__bfloat162float(a_bf[m]) * gate); + } + } + out_vec[v] = o_u4; + } + return; + } + for (int64_t column = threadIdx.x; column < output_width; column += kProjectedGateThreads) { + const int64_t index = row_base + column; + __nv_bfloat16 value = __float2bfloat16_rn(__bfloat162float(attention[index]) * gate); + if (residual != nullptr) { + value = __float2bfloat16_rn(__bfloat162float(value) + __bfloat162float(residual[index])); + } + output[index] = value; + } +} + template __device__ __forceinline__ T RoundedMulAdd(T x, T y, T gate); @@ -141,6 +232,24 @@ INSTANTIATE_GATED_ADD(__nv_bfloat16) #undef INSTANTIATE_GATED_ADD +Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, + const uint16_t* activation, const uint16_t* weight, + const uint16_t* residual, uint16_t* output, + int64_t rows, int64_t input_width, int64_t output_width, + bool sigmoid_in_fp32) { + ORT_RETURN_IF_NOT(rows >= 0 && rows <= std::numeric_limits::max() && + input_width > 0 && output_width > 0, + "GatedAdd projected gate dimensions exceed CUDA limits"); + if (rows == 0) return Status::OK(); + ProjectedSigmoidGateMulKernel<<(rows), kProjectedGateThreads, 0, stream>>>( + reinterpret_cast(attention), + reinterpret_cast(activation), + reinterpret_cast(weight), + reinterpret_cast(residual), reinterpret_cast<__nv_bfloat16*>(output), + input_width, output_width, sigmoid_in_fp32); + return CUDA_CALL(cudaGetLastError()); +} + } // namespace cuda } // namespace contrib } // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h index 9df24bd846bb2..903dd763e5c84 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h @@ -20,6 +20,17 @@ Status LaunchGatedAddKernel(cudaStream_t stream, int64_t count, int64_t hidden_size); +Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, + const uint16_t* attention, + const uint16_t* activation, + const uint16_t* weight, + const uint16_t* residual, + uint16_t* output, + int64_t rows, + int64_t input_width, + int64_t output_width, + bool sigmoid_in_fp32); + } // namespace cuda } // namespace contrib } // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cuda/bert/scaled_silu_impl.cu b/onnxruntime/contrib_ops/cuda/bert/scaled_silu_impl.cu index 249c44d70ce11..f05bc281b975a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/scaled_silu_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/scaled_silu_impl.cu @@ -40,43 +40,12 @@ __global__ void ScaledSiLUKernel(const T* x, const void* scale, int scale_type, y[i] = from_float(to_float(z) * to_float(sigmoid)); } -template -__global__ void ScaledSiLUVec8Kernel(const uint4* x, const void* scale, int scale_type, - uint4* y, int vec_count, float alpha) { - const int v = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (v >= vec_count) return; - const float scale_value = scale == nullptr ? alpha : LoadScale(scale, scale_type); - const uint4 x_u4 = x[v]; - uint4 y_u4; - const T* x_t = reinterpret_cast(&x_u4); - T* y_t = reinterpret_cast(&y_u4); -#pragma unroll - for (int m = 0; m < 8; ++m) { - const T z = from_float(to_float(x_t[m]) * scale_value); - const T sigmoid = from_float(Sigmoid(to_float(z))); - y_t[m] = from_float(to_float(z) * to_float(sigmoid)); - } - y[v] = y_u4; -} - } // namespace template Status LaunchScaledSiLU(cudaStream_t stream, const T* x, const void* scale, int scale_type, T* y, int64_t count, float alpha) { if (count == 0) return Status::OK(); - if constexpr (sizeof(T) == 2) { - if ((count & 7) == 0 && (count >> 3) <= std::numeric_limits::max() && - (reinterpret_cast(x) & 15) == 0 && - (reinterpret_cast(y) & 15) == 0) { - const int vec_count = static_cast(count >> 3); - const int vec_blocks = (vec_count - 1) / kThreads + 1; - ScaledSiLUVec8Kernel<<>>( - reinterpret_cast(x), scale, scale_type, - reinterpret_cast(y), vec_count, alpha); - return CUDA_CALL(cudaGetLastError()); - } - } const int64_t block_count = (count - 1) / kThreads + 1; ORT_RETURN_IF_NOT(block_count <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); ScaledSiLUKernel<<(block_count), kThreads, 0, stream>>>(x, scale, scale_type, y, count, alpha); diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index c75488f452caa..e8f435f99b4d7 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -170,7 +170,6 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, LinearAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, GatedDeltaNet); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, GatedDeltaNet); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, GatedDeltaNet); -class CUDA_MS_OP_CLASS_NAME(1, SigmoidGateMul); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, LinearAttentionGate); @@ -538,7 +537,6 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 9739da6491154..bf8f0ea932e86 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -5721,48 +5721,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } })); -constexpr const char* FusedTensorOps_ver1_doc = R"DOC( -CUDA contrib operators for fused tensor operations used by transformer models. -)DOC"; - -ONNX_MS_OPERATOR_SET_SCHEMA( - SigmoidGateMul, 1, - OpSchema() - .SetDoc(FusedTensorOps_ver1_doc) - .Attr("head_dim", "Dimension of each attention head.", AttributeProto::INT, static_cast(0)) - .Attr("num_heads", "Number of attention heads.", AttributeProto::INT, static_cast(0)) - .Attr("q_stride", "Per-head stride in the packed projection.", AttributeProto::INT, static_cast(0)) - .Attr("tok_stride", "Per-token stride in the packed projection.", AttributeProto::INT, static_cast(0)) - .Attr("q_offset", "Offset of the query projection within each token.", AttributeProto::INT, static_cast(0)) - .Attr("broadcast_gate", - "When one, projection has the attention shape with final dimension one. " - "Packed layout attributes must be zero. CUDA only.", - AttributeProto::INT, static_cast(0)) - .Attr("sigmoid_in_fp32", - "One computes sigmoid in FP32 before BF16 conversion. Zero preserves " - "the input-type intermediate arithmetic of CUDA Sigmoid, including subtraction on negative inputs.", - AttributeProto::INT, static_cast(1)) - .Attr("silu_gate", - "When one, projection and attention have identical shapes. Multiply the " - "BF16-rounded sigmoid by projection and round to BF16 before multiplying attention. " - "Cannot be combined with broadcast_gate or packed layout attributes.", - AttributeProto::INT, static_cast(0)) - .Input(0, "attention", "BF16 attention values.", "T") - .Input(1, "projection", "BF16 packed projection containing gate values.", "T") - .Input(2, "residual", "Optional BF16 residual with the attention shape.", "T", OpSchema::Optional) - .Input(3, "prefill_numerics", "Optional CPU int32 scalar 0/1.", "TI", OpSchema::Optional) - .Input(4, "gate_weight", "Optional BF16 projected-gate weight [input_width,1], requires broadcast_gate.", "T", OpSchema::Optional) - .Output(0, "output", - "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 " - "before optionally adding residual and rounding again.", - "T") - .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and output use BF16.") - .TypeConstraint("TI", {"tensor(int32)"}, "CPU numerical control.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - if (hasInputShape(ctx, 0)) updateOutputShape(ctx, 0, getInputShape(ctx, 0)); - })); - constexpr const char* LinearAttentionGate_ver1_doc = R"DOC( Fuses the gate projections that feed LinearAttention's gated-delta recurrence: diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 53842f1666c0f..bec1ef0fde030 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -97,7 +97,6 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPostMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionLowRankMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNet); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNetStateReplay); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SigmoidGateMul); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ConvStateReplay); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedRMSNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedAdd); @@ -234,7 +233,6 @@ 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/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc deleted file mode 100644 index 1bf2b9f39b185..0000000000000 --- a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc +++ /dev/null @@ -1,302 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include -#include -#include -#include -#include -#include - -#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(FusedTensorOpsTest, ArgMaxLastAxisWithoutKeepdims) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - - OpTester test("ArgMax", 13); - test.AddAttribute("axis", -1); - test.AddAttribute("keepdims", 0); - test.AddInput("logits", {2, 4}, {1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f}); - test.AddOutput("indices", {2}, {1, 0}); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - -TEST(FusedTensorOpsTest, ArgMaxLastAxisWithKeepdims) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - - OpTester test("ArgMax", 13); - test.AddAttribute("axis", -1); - test.AddAttribute("keepdims", 1); - test.AddInput("logits", {2, 4}, - ToTensorType({1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f})); - test.AddOutput("indices", {2, 1}, {1, 0}); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - -TEST(FusedTensorOpsTest, ArgMaxLastAxisPartitionedVocabulary) { - constexpr int64_t width = 248321; - for (bool keepdims : {false, true}) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - std::vector values(3 * width, -std::numeric_limits::infinity()); - values[4095] = values[4096] = values[width - 1] = 5.0f; - values[2 * width - 1] = 7.0f; - OpTester test("ArgMax", 13); - test.AddAttribute("axis", -1); - test.AddAttribute("keepdims", keepdims ? 1 : 0); - test.AddInput("logits", {3, width}, ToBFloat16(values)); - test.AddOutput("indices", keepdims ? std::vector{3, 1} : std::vector{3}, - {4095, width - 1, 0}); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); - } -} - -template -void RunArgMaxNaNRows() { - for (int64_t width : {1, 4096, 4097, 12289}) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - const float nan = std::numeric_limits::quiet_NaN(); - std::vector values(4 * width, nan); - values[width] = -std::numeric_limits::infinity(); - values[2 * width - 1] = -std::numeric_limits::infinity(); - values[2 * width] = values[3 * width - 1] = 5.0f; - values[3 * width] = -1.0f; - values[4 * width - 1] = 7.0f; - OpTester test("ArgMax", 13); - test.AddAttribute("axis", -1); - test.AddAttribute("keepdims", 0); - test.AddInput("logits", {4, width}, ToTensorType(values)); - test.AddOutput("indices", {4}, {0, 0, 0, width - 1}); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); - } -} - -TEST(FusedTensorOpsTest, ArgMaxLastAxisNaNSingleAndMultiplePartitions) { - RunArgMaxNaNRows(); - RunArgMaxNaNRows(); - RunArgMaxNaNRows(); -} - -TEST(FusedTensorOpsTest, SigmoidGateMul) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - - OpTester test("SigmoidGateMul", 1, kMSDomain); - test.AddAttribute("head_dim", 2); - test.AddAttribute("num_heads", 1); - test.AddAttribute("q_stride", 4); - test.AddAttribute("tok_stride", 4); - test.AddAttribute("q_offset", 0); - test.AddInput("attention", {1, 1, 2}, ToTensorType({2.0f, 4.0f})); - test.AddInput("projection", {1, 4}, ToTensorType({1.0f, 2.0f, 0.0f, 0.0f})); - test.AddOutput("output", {1, 1, 2}, ToTensorType({1.0f, 2.0f})); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - -TEST(FusedTensorOpsTest, SigmoidGateMulCheckedProjectionExtent) { - for (int mode = 0; mode < 4; ++mode) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - const bool empty = mode == 3; - OpTester test("SigmoidGateMul", 1, kMSDomain); - test.AddAttribute("head_dim", 1); - test.AddAttribute("num_heads", 1); - test.AddAttribute("q_stride", 2); - test.AddAttribute("tok_stride", mode == 1 ? 1 : std::numeric_limits::max()); - test.AddAttribute("q_offset", mode == 1 ? std::numeric_limits::max() : 0); - const int64_t tokens = empty ? 0 : (mode == 2 ? 3 : 2); - test.AddInput("attention", {tokens, 1}, std::vector(tokens)); - test.AddInput("projection", {0}, {}); - test.AddOutput("output", {tokens, 1}, std::vector(tokens)); - test.Run(empty ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, - empty ? "" : "projection extent overflows int64", {}, nullptr, &providers); - } -} - -TEST(FusedTensorOpsTest, SigmoidGateMulPaddedHeadStride) { - for (int64_t tokens : {1, 2}) { - for (bool truncated : {false, true}) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - OpTester test("SigmoidGateMul", 1, kMSDomain); - test.AddAttribute("head_dim", 1); - test.AddAttribute("num_heads", 2); - test.AddAttribute("q_stride", 4); - test.AddAttribute("tok_stride", 6); - const int64_t extent = tokens * 6 - (truncated ? 1 : 0); - test.AddInput("attention", {tokens, 2}, - tokens == 1 ? ToBFloat16({2, 4}) : ToBFloat16({2, 4, 6, 8})); - test.AddInput("projection", {extent}, std::vector(extent)); - test.AddOutput("output", {tokens, 2}, - tokens == 1 ? ToBFloat16({1, 2}) : ToBFloat16({1, 2, 3, 4})); - test.Run(truncated ? OpTester::ExpectResult::kExpectFailure : OpTester::ExpectResult::kExpectSuccess, - truncated ? "projection is too small" : "", {}, nullptr, &providers); - } - } -} - -TEST(FusedTensorOpsTest, SigmoidGateMulRejectsOutOfRangeLayoutAttributes) { -#if defined(ORT_NO_EXCEPTIONS) - GTEST_SKIP() << "Checked narrowing failures abort when exceptions are disabled."; -#endif - for (const char* attribute : {"head_dim", "num_heads", "q_stride"}) { - for (int64_t value : {int64_t{4294967300}, int64_t{-4294967292}}) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - SCOPED_TRACE(attribute); - SCOPED_TRACE(value); - OpTester test("SigmoidGateMul", 1, kMSDomain); - test.AddAttribute("head_dim", std::string(attribute) == "head_dim" ? value : 2); - test.AddAttribute("num_heads", std::string(attribute) == "num_heads" ? value : 2); - test.AddAttribute("q_stride", std::string(attribute) == "q_stride" ? value : 4); - test.AddAttribute("tok_stride", 8); - test.AddInput("attention", {1, 4}, std::vector(4)); - test.AddInput("projection", {8}, std::vector(8)); - test.AddOutput("output", {1, 4}, std::vector(4)); - test.Run(OpTester::ExpectResult::kExpectFailure, "narrowing_error", {}, nullptr, &providers); - } - } -} - -void RunSigmoidGateMulBroadcastResidualRounding(bool sigmoid_in_fp32, bool silu_gate = false) { - for (bool residual : {false, true}) { - for (int64_t width : {1, 257, 2560}) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - const std::vector base_gates = ToBFloat16({-100.0f, -5.5f, -0.0f, 0.3f, 5.5f, 100.0f}); - std::vector gates(silu_gate ? 6 * width : 6); - for (size_t i = 0; i < gates.size(); ++i) gates[i] = base_gates[i % 6]; - std::vector values(6 * width), residuals(6 * width), expected(6 * width); - for (size_t i = 0; i < values.size(); ++i) { - values[i] = BFloat16(static_cast(static_cast(i % 23) - 11) * 0.27f); - residuals[i] = BFloat16(static_cast(static_cast(i % 17) - 8) * 0.13f); - const float gate = gates[silu_gate ? i : i / width].ToFloat(); - const float exp_value = std::exp(gate >= 0 ? -gate : gate); - const float sigmoid = gate >= 0 ? 1.0f / (1.0f + exp_value) : exp_value / (1.0f + exp_value); - const BFloat16 denominator(1.0f + BFloat16(exp_value).ToFloat()); - const BFloat16 reciprocal(1.0f / denominator.ToFloat()); - const BFloat16 rounded_sigmoid = sigmoid_in_fp32 - ? BFloat16(sigmoid) - : (gate > 0 ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat())); - const BFloat16 activation = silu_gate ? BFloat16(gate * rounded_sigmoid.ToFloat()) : rounded_sigmoid; - const BFloat16 product(values[i].ToFloat() * activation.ToFloat()); - expected[i] = residual ? BFloat16(product.ToFloat() + residuals[i].ToFloat()) : product; - } - OpTester test("SigmoidGateMul", 1, kMSDomain); - test.AddAttribute(silu_gate ? "silu_gate" : "broadcast_gate", 1); - test.AddAttribute("sigmoid_in_fp32", sigmoid_in_fp32 ? 1 : 0); - test.AddInput("attention", {2, 3, width}, values); - test.AddInput("projection", {2, 3, silu_gate ? width : 1}, gates); - if (residual) test.AddInput("residual", {2, 3, width}, residuals); - test.AddOutput("output", {2, 3, width}, expected); - test.SetOutputTolerance(0.0f, 0.0f); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); - } - } -} - -TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastResidualRounding) { - RunSigmoidGateMulBroadcastResidualRounding(true); -} - -TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastInputPrecision) { - RunSigmoidGateMulBroadcastResidualRounding(false); -} - -TEST(FusedTensorOpsTest, SigmoidGateMulSiluRounding) { - RunSigmoidGateMulBroadcastResidualRounding(true, true); - RunSigmoidGateMulBroadcastResidualRounding(false, true); -} - -TEST(FusedTensorOpsTest, SigmoidGateMulRejectsBroadcastAndResidualShapes) { - for (bool bad_residual : {false, true}) { - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - OpTester test("SigmoidGateMul", 1, kMSDomain); - test.AddAttribute("broadcast_gate", 1); - test.AddInput("attention", {2, 3}, ToBFloat16(std::vector(6, 1.0f))); - test.AddInput("projection", bad_residual ? std::vector{2, 1} : std::vector{1, 2}, - ToBFloat16({1.0f, 1.0f})); - if (bad_residual) test.AddInput("residual", {1}, ToBFloat16({1.0f})); - test.AddOutput("output", {2, 3}, ToBFloat16(std::vector(6, 0.0f))); - test.Run(OpTester::ExpectResult::kExpectFailure, - bad_residual ? "residual must have the same shape" : "singleton final dimension", - {}, nullptr, &providers); - } -} - -TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastEmptyAndInvalidAttributes) { - for (int mode = 0; mode < 4; ++mode) { -#if defined(ORT_NO_EXCEPTIONS) - if (mode != 0) continue; -#endif - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - OpTester test("SigmoidGateMul", 1, kMSDomain); - test.AddAttribute("broadcast_gate", mode == 1 ? 2 : 1); - if (mode == 2) test.AddAttribute("sigmoid_in_fp32", 2); - if (mode == 3) test.AddAttribute("head_dim", 2); - test.AddInput("attention", {0, 3}, {}); - test.AddInput("projection", {0, 1}, {}); - test.AddInput("residual", {0, 3}, {}); - test.AddOutput("output", {0, 3}, {}); - const std::vector errors = {"", "broadcast_gate must be zero or one", - "sigmoid_in_fp32 must be zero or one", - "cannot be combined with packed layout attributes"}; - test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, - errors[mode], {}, nullptr, &providers); - } -} - -TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) { - for (int mode = 0; mode < 5; ++mode) { -#if defined(ORT_NO_EXCEPTIONS) - if (mode >= 1 && mode <= 3) continue; -#endif - auto providers = CudaProvider(); - if (providers.empty()) GTEST_SKIP(); - OpTester test("SigmoidGateMul", 1, kMSDomain); - test.AddAttribute("silu_gate", mode == 1 ? 2 : 1); - if (mode == 2) test.AddAttribute("broadcast_gate", 1); - if (mode == 3) test.AddAttribute("head_dim", 2); - test.AddInput("attention", {0, 3}, {}); - test.AddInput("projection", {0, mode == 4 ? 1 : 3}, {}); - test.AddOutput("output", {0, 3}, {}); - const std::vector errors = {"", "silu_gate must be zero or one", "are exclusive", - "cannot be combined with packed layout attributes", - "silu gate must have the same shape"}; - test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, - errors[mode], {}, nullptr, &providers); - } -} - -} // namespace -} // namespace onnxruntime::test diff --git a/onnxruntime/test/contrib_ops/gated_add_op_test.cc b/onnxruntime/test/contrib_ops/gated_add_op_test.cc index d3a978c076a5a..98272410b9dc6 100644 --- a/onnxruntime/test/contrib_ops/gated_add_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_add_op_test.cc @@ -179,5 +179,81 @@ TEST(ContribOpGatedAddTest, MismatchedYShape) { } } +TEST(ContribOpGatedAddTest, ProjectedGateBFloat16) { + if (!CudaHasBF16Support()) { + GTEST_SKIP() << "bfloat16 requires compute capability 8.0 or later"; + } + for (bool sigmoid_in_fp32 : {false, true}) { + for (int64_t width : {8, 2560}) { + constexpr int64_t kRows = 4; + constexpr int64_t kInputWidth = 16; + std::vector x(kRows * width); + std::vector y(kRows * width); + std::vector act(kRows * kInputWidth); + std::vector weight(kInputWidth); + std::vector expected(kRows * width); + for (int64_t k = 0; k < kInputWidth; ++k) { + weight[k] = BFloat16(static_cast((k % 5) - 2) * 0.25f); + } + for (int64_t r = 0; r < kRows; ++r) { + float dot = 0.0f; + for (int64_t k = 0; k < kInputWidth; ++k) { + act[r * kInputWidth + k] = BFloat16(static_cast(((r + k) % 7) - 3) * 0.3f); + dot = std::fma(act[r * kInputWidth + k].ToFloat(), weight[k].ToFloat(), dot); + } + const float gate = BFloat16(dot).ToFloat(); + const float exp_value = std::exp(gate >= 0.0f ? -gate : gate); + const float sigmoid = gate >= 0.0f ? 1.0f / (1.0f + exp_value) : exp_value / (1.0f + exp_value); + const BFloat16 denominator(1.0f + BFloat16(exp_value).ToFloat()); + const BFloat16 reciprocal(1.0f / denominator.ToFloat()); + const BFloat16 rounded_sigmoid = + sigmoid_in_fp32 ? BFloat16(sigmoid) + : (gate > 0.0f ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat())); + for (int64_t c = 0; c < width; ++c) { + const int64_t idx = r * width + c; + x[idx] = BFloat16(static_cast((idx % 17) - 8) * 0.13f); + y[idx] = BFloat16(static_cast((idx % 11) - 5) * 0.37f); + const BFloat16 product(y[idx].ToFloat() * rounded_sigmoid.ToFloat()); + expected[idx] = BFloat16(x[idx].ToFloat() + product.ToFloat()); + } + } + std::vector> providers; + if (auto cuda_ep = DefaultCudaExecutionProvider()) { + providers.push_back(std::move(cuda_ep)); + } + if (providers.empty()) GTEST_SKIP(); + OpTester tester("GatedAdd", 1, onnxruntime::kMSDomain); + tester.AddAttribute("sigmoid_in_fp32", sigmoid_in_fp32 ? 1 : 0); + tester.AddInput("X", {1, kRows, width}, x); + tester.AddInput("Y", {1, kRows, width}, y); + tester.AddInput("gate", {1, kRows, kInputWidth}, act); + tester.AddInput("gate_weight", {kInputWidth, 1}, weight); + tester.AddOutput("output", {1, kRows, width}, expected); + tester.SetOutputTolerance(0.0f, 0.0f); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + } +} + +TEST(ContribOpGatedAddTest, ArgMaxBFloat16) { + if (!CudaHasBF16Support()) { + GTEST_SKIP() << "bfloat16 requires compute capability 8.0 or later"; + } + for (bool keepdims : {false, true}) { + std::vector> providers; + if (auto cuda_ep = DefaultCudaExecutionProvider()) { + providers.push_back(std::move(cuda_ep)); + } + if (providers.empty()) GTEST_SKIP(); + OpTester test("ArgMax", 13); + test.AddAttribute("axis", -1); + test.AddAttribute("keepdims", keepdims ? 1 : 0); + test.AddInput("logits", {2, 4}, + ToTensorType({1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f})); + test.AddOutput("indices", keepdims ? std::vector{2, 1} : std::vector{2}, {1, 0}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } +} + } // namespace test } // namespace onnxruntime From a9b22e4c101fc773960ea52d0cc54e292bc443b0 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 06:07:48 +0000 Subject: [PATCH 16/19] Address PR #33470 review feedback for GatedAdd and BFloat16 ArgMax/ArgMin - Explicitly reject gate_weight on CPU and WebGPU GatedAdd kernels and validate sigmoid_in_fp32 attribute - Add K=1 projected gate regression test across EPs and odd-width / K>256 BFloat16 projected gate coverage - Add BFloat16 single/multi-block, tie, and special-value cases to reduction_functions_test.cc and non-last-axis ArgMax/ArgMin tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cpu/bert/gated_add.cc | 4 + onnxruntime/contrib_ops/cpu/bert/gated_add.h | 5 +- .../contrib_ops/webgpu/bert/gated_add.cc | 4 + .../contrib_ops/webgpu/bert/gated_add.h | 5 +- .../test/contrib_ops/gated_add_op_test.cc | 156 +++++++++++++++--- .../cpu/reduction/reduction_ops_test.cc | 48 ++++++ .../test_cases/reduction_functions_test.cc | 8 + 7 files changed, 207 insertions(+), 23 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/gated_add.cc b/onnxruntime/contrib_ops/cpu/bert/gated_add.cc index 48b643da3bc20..8c79885a3ee33 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gated_add.cc +++ b/onnxruntime/contrib_ops/cpu/bert/gated_add.cc @@ -46,6 +46,10 @@ Status GatedAdd::Compute(OpKernelContext* context) const { const Tensor* x = context->Input(0); const Tensor* y = context->Input(1); const Tensor* gate = context->Input(2); + const Tensor* gate_weight = context->Input(3); + if (gate_weight != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "GatedAdd gate_weight currently supports BFloat16 only"); + } const TensorShape& shape = x->Shape(); ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1, "X must have rank >= 1"); diff --git a/onnxruntime/contrib_ops/cpu/bert/gated_add.h b/onnxruntime/contrib_ops/cpu/bert/gated_add.h index b0008e29b3cc1..4de8184f56d8a 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gated_add.h +++ b/onnxruntime/contrib_ops/cpu/bert/gated_add.h @@ -13,7 +13,10 @@ namespace contrib { template class GatedAdd final : public OpKernel { public: - explicit GatedAdd(const OpKernelInfo& info) : OpKernel(info) {} + explicit GatedAdd(const OpKernelInfo& info) : OpKernel(info) { + const int64_t fp32 = info.GetAttrOrDefault("sigmoid_in_fp32", 1); + ORT_ENFORCE(fp32 == 0 || fp32 == 1, "GatedAdd sigmoid_in_fp32 must be zero or one"); + } Status Compute(OpKernelContext* context) const override; }; diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_add.cc b/onnxruntime/contrib_ops/webgpu/bert/gated_add.cc index 1202fe55857a6..2e24fc8bcdc6c 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/gated_add.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_add.cc @@ -40,6 +40,10 @@ Status GatedAdd::ComputeInternal(ComputeContext& context) const { const auto* x = context.Input(0); const auto* y = context.Input(1); const auto* gate = context.Input(2); + const auto* gate_weight = context.Input(3); + if (gate_weight != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "GatedAdd gate_weight currently supports BFloat16 only"); + } const TensorShape& shape = x->Shape(); ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1, "X must have rank >= 1"); diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_add.h b/onnxruntime/contrib_ops/webgpu/bert/gated_add.h index 1bb06237a39ff..e445c2fafdb6a 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/gated_add.h +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_add.h @@ -24,7 +24,10 @@ class GatedAddProgram final : public Program { class GatedAdd final : public WebGpuKernel { public: - GatedAdd(const OpKernelInfo& info) : WebGpuKernel(info) {} + GatedAdd(const OpKernelInfo& info) : WebGpuKernel(info) { + const int64_t fp32 = info.GetAttrOrDefault("sigmoid_in_fp32", 1); + ORT_ENFORCE(fp32 == 0 || fp32 == 1, "GatedAdd sigmoid_in_fp32 must be zero or one"); + } Status ComputeInternal(ComputeContext& context) const override; }; diff --git a/onnxruntime/test/contrib_ops/gated_add_op_test.cc b/onnxruntime/test/contrib_ops/gated_add_op_test.cc index 98272410b9dc6..7f38305454a81 100644 --- a/onnxruntime/test/contrib_ops/gated_add_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_add_op_test.cc @@ -183,23 +183,28 @@ TEST(ContribOpGatedAddTest, ProjectedGateBFloat16) { if (!CudaHasBF16Support()) { GTEST_SKIP() << "bfloat16 requires compute capability 8.0 or later"; } + struct ProjectionCase { + int64_t width; + int64_t input_width; + }; for (bool sigmoid_in_fp32 : {false, true}) { - for (int64_t width : {8, 2560}) { + for (const ProjectionCase& test_case : {ProjectionCase{8, 16}, ProjectionCase{2560, 16}, ProjectionCase{259, 300}}) { constexpr int64_t kRows = 4; - constexpr int64_t kInputWidth = 16; + const int64_t width = test_case.width; + const int64_t input_width = test_case.input_width; std::vector x(kRows * width); std::vector y(kRows * width); - std::vector act(kRows * kInputWidth); - std::vector weight(kInputWidth); + std::vector act(kRows * input_width); + std::vector weight(input_width); std::vector expected(kRows * width); - for (int64_t k = 0; k < kInputWidth; ++k) { + for (int64_t k = 0; k < input_width; ++k) { weight[k] = BFloat16(static_cast((k % 5) - 2) * 0.25f); } for (int64_t r = 0; r < kRows; ++r) { float dot = 0.0f; - for (int64_t k = 0; k < kInputWidth; ++k) { - act[r * kInputWidth + k] = BFloat16(static_cast(((r + k) % 7) - 3) * 0.3f); - dot = std::fma(act[r * kInputWidth + k].ToFloat(), weight[k].ToFloat(), dot); + for (int64_t k = 0; k < input_width; ++k) { + act[r * input_width + k] = BFloat16(static_cast(((r + k) % 7) - 3) * 0.125f); + dot = std::fma(act[r * input_width + k].ToFloat(), weight[k].ToFloat(), dot); } const float gate = BFloat16(dot).ToFloat(); const float exp_value = std::exp(gate >= 0.0f ? -gate : gate); @@ -226,8 +231,8 @@ TEST(ContribOpGatedAddTest, ProjectedGateBFloat16) { tester.AddAttribute("sigmoid_in_fp32", sigmoid_in_fp32 ? 1 : 0); tester.AddInput("X", {1, kRows, width}, x); tester.AddInput("Y", {1, kRows, width}, y); - tester.AddInput("gate", {1, kRows, kInputWidth}, act); - tester.AddInput("gate_weight", {kInputWidth, 1}, weight); + tester.AddInput("gate", {1, kRows, input_width}, act); + tester.AddInput("gate_weight", {input_width, 1}, weight); tester.AddOutput("output", {1, kRows, width}, expected); tester.SetOutputTolerance(0.0f, 0.0f); tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); @@ -235,23 +240,132 @@ TEST(ContribOpGatedAddTest, ProjectedGateBFloat16) { } } +TEST(ContribOpGatedAddTest, ProjectedGateK1Regression) { + for (auto& ep : AvailableGatedOpExecutionProviders()) { + SCOPED_TRACE("EP (float): " + ep->Type()); + OpTester tester("GatedAdd", 1, onnxruntime::kMSDomain); + tester.AddInput("X", {1, 1}, {0.0f}); + tester.AddInput("Y", {1, 1}, {2.0f}); + tester.AddInput("gate", {1, 1}, {0.0f}); + tester.AddInput("gate_weight", {1, 1}, {1.0f}); + tester.AddOutput("output", {1, 1}, {1.0f}); + + std::vector> providers; + providers.push_back(std::move(ep)); + tester.Run(OpTester::ExpectResult::kExpectFailure, "GatedAdd gate_weight currently supports BFloat16 only", + {}, nullptr, &providers); + } + + for (auto& ep : AvailableGatedOpExecutionProviders()) { + SCOPED_TRACE("EP (MLFloat16): " + ep->Type()); + OpTester tester("GatedAdd", 1, onnxruntime::kMSDomain); + tester.AddInput("X", {1, 1}, ToTensorType({0.0f})); + tester.AddInput("Y", {1, 1}, ToTensorType({2.0f})); + tester.AddInput("gate", {1, 1}, ToTensorType({0.0f})); + tester.AddInput("gate_weight", {1, 1}, ToTensorType({1.0f})); + tester.AddOutput("output", {1, 1}, ToTensorType({1.0f})); + + std::vector> providers; + providers.push_back(std::move(ep)); + tester.Run(OpTester::ExpectResult::kExpectFailure, "GatedAdd gate_weight currently supports BFloat16 only", + {}, nullptr, &providers); + } + + if (CudaHasBF16Support()) { + std::vector> providers; + if (auto cuda_ep = DefaultCudaExecutionProvider()) { + providers.push_back(std::move(cuda_ep)); + } + if (!providers.empty()) { + OpTester tester("GatedAdd", 1, onnxruntime::kMSDomain); + tester.AddInput("X", {1, 1}, ToTensorType({0.0f})); + tester.AddInput("Y", {1, 1}, ToTensorType({2.0f})); + tester.AddInput("gate", {1, 1}, ToTensorType({0.0f})); + tester.AddInput("gate_weight", {1, 1}, ToTensorType({1.0f})); + tester.AddOutput("output", {1, 1}, ToTensorType({1.0f})); + tester.SetOutputTolerance(0.0f, 0.0f); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + } +} + TEST(ContribOpGatedAddTest, ArgMaxBFloat16) { if (!CudaHasBF16Support()) { GTEST_SKIP() << "bfloat16 requires compute capability 8.0 or later"; } for (bool keepdims : {false, true}) { - std::vector> providers; - if (auto cuda_ep = DefaultCudaExecutionProvider()) { - providers.push_back(std::move(cuda_ep)); + for (const char* op : {"ArgMax", "ArgMin"}) { + const bool is_max = std::string(op) == "ArgMax"; + std::vector> providers; + if (auto cuda_ep = DefaultCudaExecutionProvider()) { + providers.push_back(std::move(cuda_ep)); + } + if (providers.empty()) GTEST_SKIP(); + OpTester test(op, 13); + test.AddAttribute("axis", -1); + test.AddAttribute("keepdims", keepdims ? 1 : 0); + test.AddInput("logits", {2, 4}, + ToTensorType({1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f})); + test.AddOutput("indices", keepdims ? std::vector{2, 1} : std::vector{2}, + is_max ? std::vector{1, 0} : std::vector{3, 3}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + + for (int64_t cols : {int64_t{1024}, int64_t{40000}}) { + constexpr int64_t kRows = 2; + const std::vector max_indices{17, cols - 5}; + const std::vector min_indices{29, cols / 2}; + std::vector logits(static_cast(kRows * cols)); + for (int64_t r = 0; r < kRows; ++r) { + for (int64_t c = 0; c < cols; ++c) { + logits[static_cast(r * cols + c)] = static_cast((c + r) % 31) * 0.5f; + } + logits[static_cast(r * cols + max_indices[static_cast(r)])] = 256.0f; + logits[static_cast(r * cols + cols - 1)] = 256.0f; + logits[static_cast(r * cols + min_indices[static_cast(r)])] = -256.0f; + logits[static_cast(r * cols + cols - 2)] = -256.0f; + } + for (const char* op : {"ArgMax", "ArgMin"}) { + const bool is_max = std::string(op) == "ArgMax"; + std::vector> providers; + if (auto cuda_ep = DefaultCudaExecutionProvider()) { + providers.push_back(std::move(cuda_ep)); + } + if (providers.empty()) GTEST_SKIP(); + OpTester test(op, 13); + test.AddAttribute("axis", -1); + test.AddAttribute("keepdims", keepdims ? 1 : 0); + test.AddInput("logits", {kRows, cols}, ToTensorType(logits)); + test.AddOutput("indices", keepdims ? std::vector{kRows, 1} : std::vector{kRows}, + is_max ? max_indices : min_indices); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + } + + for (const char* op : {"ArgMax", "ArgMin"}) { + const bool is_max = std::string(op) == "ArgMax"; + std::vector> providers; + if (auto cuda_ep = DefaultCudaExecutionProvider()) { + providers.push_back(std::move(cuda_ep)); + } + if (providers.empty()) GTEST_SKIP(); + OpTester test(op, 13); + test.AddAttribute("axis", 1); + test.AddAttribute("keepdims", keepdims ? 1 : 0); + test.AddInput("logits", {3, 2, 2}, + ToTensorType({1.0f, 4.0f, + 3.0f, 2.0f, + + 8.0f, 5.0f, + 7.0f, 6.0f, + + 9.0f, 12.0f, + 11.0f, 10.0f})); + test.AddOutput("indices", keepdims ? std::vector{3, 1, 2} : std::vector{3, 2}, + is_max ? std::vector{1, 0, 0, 1, 1, 0} + : std::vector{0, 1, 1, 0, 0, 1}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } - if (providers.empty()) GTEST_SKIP(); - OpTester test("ArgMax", 13); - test.AddAttribute("axis", -1); - test.AddAttribute("keepdims", keepdims ? 1 : 0); - test.AddInput("logits", {2, 4}, - ToTensorType({1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f})); - test.AddOutput("indices", keepdims ? std::vector{2, 1} : std::vector{2}, {1, 0}); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } } diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc index 03a582af013a5..228a11e633107 100644 --- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc +++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc @@ -6983,6 +6983,54 @@ TEST(ReductionOpTest, ReduceMin_float_Opset20_NoKeepdims_Cuda) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +TEST(ReductionOpTest, ArgMax_bfloat16_non_last_axis_Cuda) { + for (int64_t keepdims : {int64_t{0}, int64_t{1}}) { + OpTester test("ArgMax", 13); + test.AddAttribute("axis", int64_t{1}); + test.AddAttribute("keepdims", keepdims); + test.AddInput("data", {3, 2, 2}, + MakeBFloat16({1.0f, 4.0f, + 3.0f, 2.0f, + + 8.0f, 5.0f, + 7.0f, 6.0f, + + 9.0f, 12.0f, + 11.0f, 10.0f})); + test.AddOutput("reduced", keepdims ? std::vector{3, 1, 2} : std::vector{3, 2}, + {1, 0, + 0, 1, + 1, 0}); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} + +TEST(ReductionOpTest, ArgMin_bfloat16_non_last_axis_Cuda) { + for (int64_t keepdims : {int64_t{0}, int64_t{1}}) { + OpTester test("ArgMin", 13); + test.AddAttribute("axis", int64_t{1}); + test.AddAttribute("keepdims", keepdims); + test.AddInput("data", {3, 2, 2}, + MakeBFloat16({1.0f, 4.0f, + 3.0f, 2.0f, + + 8.0f, 5.0f, + 7.0f, 6.0f, + + 9.0f, 12.0f, + 11.0f, 10.0f})); + test.AddOutput("reduced", keepdims ? std::vector{3, 1, 2} : std::vector{3, 2}, + {0, 1, + 1, 0, + 0, 1}); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} + #endif // defined(USE_CUDA) // ============================================================================= diff --git a/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc b/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc index d4e5d5e923957..2f4af0288f38b 100644 --- a/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc @@ -421,6 +421,12 @@ struct ArgTestTraits { static double ToDouble(half value) { return static_cast(__half2float(value)); } }; +template <> +struct ArgTestTraits { + static BFloat16 FromFloat(float value) { return BFloat16(value); } + static double ToDouble(BFloat16 value) { return static_cast(value.ToFloat()); } +}; + // Mirrors the sequential semantics of the reduction: seed with the first // element, keep the first occurrence of the extreme value and never select a // NaN (every comparison against NaN is false, so a leading NaN wins the row). @@ -584,6 +590,7 @@ TEST(ReductionFunctionsTest, ArgMinMaxLastAxisManyRows) { TEST(ReductionFunctionsTest, ArgMinMaxLastAxisTypes) { for (int n : {64, 1024, 65536}) { TestArgMinMaxLastAxisRandom(3, n, n); + TestArgMinMaxLastAxisRandom(3, n, n); TestArgMinMaxLastAxisRandom(3, n, n); TestArgMinMaxLastAxisRandom(3, n, n); } @@ -593,6 +600,7 @@ TEST(ReductionFunctionsTest, ArgMinMaxLastAxisSpecialValues) { // 64 -> one thread per row, 1024 -> single block per row, 200000 -> many blocks per row. for (int n : {64, 1024, 200000}) { TestArgMinMaxLastAxisSpecialValues(n); + TestArgMinMaxLastAxisSpecialValues(n); TestArgMinMaxLastAxisSpecialValues(n); TestArgMinMaxLastAxisSpecialValues(n); } From 62be02ec39020e09e0dea22b585d6783cf9229b1 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 21:45:15 +0000 Subject: [PATCH 17/19] Remove sigmoid_in_fp32 attribute from GatedAdd and always compute projected gate sigmoid in FP32 - Remove sigmoid_in_fp32 attribute from com.microsoft::GatedAdd schema, docs, and kernel constructors - Always compute projected gate sigmoid in FP32 (StableSigmoid) before rounding to BFloat16 in ProjectedSigmoidGateMulKernel Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 7 -- onnxruntime/contrib_ops/cpu/bert/gated_add.h | 5 +- .../contrib_ops/cuda/bert/gated_add.cc | 2 +- onnxruntime/contrib_ops/cuda/bert/gated_add.h | 10 +-- .../contrib_ops/cuda/bert/gated_add_impl.cu | 25 ++---- .../contrib_ops/cuda/bert/gated_add_impl.h | 3 +- .../contrib_ops/webgpu/bert/gated_add.h | 5 +- .../core/graph/contrib_ops/bert_defs.cc | 3 - .../test/contrib_ops/gated_add_op_test.cc | 85 +++++++++---------- 9 files changed, 52 insertions(+), 93 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index c33e0c4eefc96..1ba3d71924fab 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2399,13 +2399,6 @@ This version of the operator has been available since version 1 of the 'com.micr This version of the operator has been available since version 1 of the 'com.microsoft' operator set. -#### Attributes - -
-
sigmoid_in_fp32 : int
-
When gate_weight is provided, compute sigmoid in FP32 (1, default) or input type (0).
-
- #### Inputs (3 - 4)
diff --git a/onnxruntime/contrib_ops/cpu/bert/gated_add.h b/onnxruntime/contrib_ops/cpu/bert/gated_add.h index 4de8184f56d8a..b0008e29b3cc1 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gated_add.h +++ b/onnxruntime/contrib_ops/cpu/bert/gated_add.h @@ -13,10 +13,7 @@ namespace contrib { template class GatedAdd final : public OpKernel { public: - explicit GatedAdd(const OpKernelInfo& info) : OpKernel(info) { - const int64_t fp32 = info.GetAttrOrDefault("sigmoid_in_fp32", 1); - ORT_ENFORCE(fp32 == 0 || fp32 == 1, "GatedAdd sigmoid_in_fp32 must be zero or one"); - } + explicit GatedAdd(const OpKernelInfo& info) : OpKernel(info) {} Status Compute(OpKernelContext* context) const override; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add.cc b/onnxruntime/contrib_ops/cuda/bert/gated_add.cc index 033a42c6f0673..3f5b3349d4311 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add.cc @@ -68,7 +68,7 @@ Status GatedAdd::ComputeInternal(OpKernelContext* context) const { reinterpret_cast(gate_weight->Data()), reinterpret_cast(x->Data()), reinterpret_cast(output->MutableData()), - shape.Size() / hidden_size, input_width, hidden_size, sigmoid_in_fp32_); + shape.Size() / hidden_size, input_width, hidden_size); } } diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add.h b/onnxruntime/contrib_ops/cuda/bert/gated_add.h index 2a48ab58769c3..bd6f68e6d7147 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add.h @@ -13,16 +13,8 @@ namespace cuda { template class GatedAdd final : public onnxruntime::cuda::CudaKernel { public: - explicit GatedAdd(const OpKernelInfo& info) - : CudaKernel(info), - sigmoid_in_fp32_(info.GetAttrOrDefault("sigmoid_in_fp32", 1) == 1) { - const int64_t fp32 = info.GetAttrOrDefault("sigmoid_in_fp32", 1); - ORT_ENFORCE(fp32 == 0 || fp32 == 1, "GatedAdd sigmoid_in_fp32 must be zero or one"); - } + explicit GatedAdd(const OpKernelInfo& info) : CudaKernel(info) {} Status ComputeInternal(OpKernelContext* context) const override; - - private: - bool sigmoid_in_fp32_{true}; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu index adac73eb5789b..2ce68dd290416 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu @@ -11,7 +11,6 @@ #include "contrib_ops/cuda/bert/gated_add_impl.h" #include "core/providers/cuda/cuda_common.h" -#include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/shared_inc/fast_divmod.h" namespace onnxruntime { @@ -30,19 +29,12 @@ __device__ __forceinline__ float StableSigmoid(float value) { return exp_value / (1.0f + exp_value); } -__device__ __forceinline__ float RoundedBf16Sigmoid(float value, bool sigmoid_in_fp32) { - if (sigmoid_in_fp32) { - return __bfloat162float(__float2bfloat16_rn(StableSigmoid(value))); - } - using namespace onnxruntime::cuda; - const BFloat16 a(value); - return static_cast(a > BFloat16(0.0f) - ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) - : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a))); +__device__ __forceinline__ float RoundedBf16Sigmoid(float value) { + return __bfloat162float(__float2bfloat16_rn(StableSigmoid(value))); } __device__ float ProjectedRowGate(const __nv_bfloat16* activation, const __nv_bfloat16* weight, - int64_t row, int64_t input_width, bool sigmoid_in_fp32) { + int64_t row, int64_t input_width) { using Reduce = cub::BlockReduce; __shared__ typename Reduce::TempStorage reduction; __shared__ float gate; @@ -54,7 +46,7 @@ __device__ float ProjectedRowGate(const __nv_bfloat16* activation, const __nv_bf const float total = Reduce(reduction).Sum(dot); if (threadIdx.x == 0) { const float value = __bfloat162float(__float2bfloat16_rn(total)); - gate = RoundedBf16Sigmoid(value, sigmoid_in_fp32); + gate = RoundedBf16Sigmoid(value); } __syncthreads(); return gate; @@ -63,9 +55,9 @@ __device__ float ProjectedRowGate(const __nv_bfloat16* activation, const __nv_bf __global__ void ProjectedSigmoidGateMulKernel( const __nv_bfloat16* attention, const __nv_bfloat16* activation, const __nv_bfloat16* weight, const __nv_bfloat16* residual, __nv_bfloat16* output, - int64_t input_width, int64_t output_width, bool sigmoid_in_fp32) { + int64_t input_width, int64_t output_width) { const int64_t row = blockIdx.x; - const float gate = ProjectedRowGate(activation, weight, row, input_width, sigmoid_in_fp32); + const float gate = ProjectedRowGate(activation, weight, row, input_width); const int64_t row_base = row * output_width; if ((output_width & 7) == 0 && (reinterpret_cast(attention + row_base) & 15) == 0 && @@ -235,8 +227,7 @@ INSTANTIATE_GATED_ADD(__nv_bfloat16) Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, const uint16_t* activation, const uint16_t* weight, const uint16_t* residual, uint16_t* output, - int64_t rows, int64_t input_width, int64_t output_width, - bool sigmoid_in_fp32) { + int64_t rows, int64_t input_width, int64_t output_width) { ORT_RETURN_IF_NOT(rows >= 0 && rows <= std::numeric_limits::max() && input_width > 0 && output_width > 0, "GatedAdd projected gate dimensions exceed CUDA limits"); @@ -246,7 +237,7 @@ Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attent reinterpret_cast(activation), reinterpret_cast(weight), reinterpret_cast(residual), reinterpret_cast<__nv_bfloat16*>(output), - input_width, output_width, sigmoid_in_fp32); + input_width, output_width); return CUDA_CALL(cudaGetLastError()); } diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h index 903dd763e5c84..e1e2103789569 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h @@ -28,8 +28,7 @@ Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, uint16_t* output, int64_t rows, int64_t input_width, - int64_t output_width, - bool sigmoid_in_fp32); + int64_t output_width); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_add.h b/onnxruntime/contrib_ops/webgpu/bert/gated_add.h index e445c2fafdb6a..1bb06237a39ff 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/gated_add.h +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_add.h @@ -24,10 +24,7 @@ class GatedAddProgram final : public Program { class GatedAdd final : public WebGpuKernel { public: - GatedAdd(const OpKernelInfo& info) : WebGpuKernel(info) { - const int64_t fp32 = info.GetAttrOrDefault("sigmoid_in_fp32", 1); - ORT_ENFORCE(fp32 == 0 || fp32 == 1, "GatedAdd sigmoid_in_fp32 must be zero or one"); - } + GatedAdd(const OpKernelInfo& info) : WebGpuKernel(info) {} Status ComputeInternal(ComputeContext& context) const override; }; diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index bf8f0ea932e86..aea6928e8c3a5 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -6118,9 +6118,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( GatedAdd, 1, OpSchema() .SetDoc(GatedAdd_ver1_doc) - .Attr("sigmoid_in_fp32", - "When gate_weight is provided, compute sigmoid in FP32 (1, default) or input type (0).", - AttributeProto::INT, static_cast(1)) .Input(0, "X", "Unscaled input with shape (..., C).", "T") .Input(1, "Y", "Input scaled by gate, with the same shape as X.", "T") .Input(2, "gate", "Per-row gate with shape (..., 1), or activation (..., K) when gate_weight is provided.", "T") diff --git a/onnxruntime/test/contrib_ops/gated_add_op_test.cc b/onnxruntime/test/contrib_ops/gated_add_op_test.cc index 7f38305454a81..24407becd066b 100644 --- a/onnxruntime/test/contrib_ops/gated_add_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_add_op_test.cc @@ -187,56 +187,49 @@ TEST(ContribOpGatedAddTest, ProjectedGateBFloat16) { int64_t width; int64_t input_width; }; - for (bool sigmoid_in_fp32 : {false, true}) { - for (const ProjectionCase& test_case : {ProjectionCase{8, 16}, ProjectionCase{2560, 16}, ProjectionCase{259, 300}}) { - constexpr int64_t kRows = 4; - const int64_t width = test_case.width; - const int64_t input_width = test_case.input_width; - std::vector x(kRows * width); - std::vector y(kRows * width); - std::vector act(kRows * input_width); - std::vector weight(input_width); - std::vector expected(kRows * width); + for (const ProjectionCase& test_case : {ProjectionCase{8, 16}, ProjectionCase{2560, 16}, ProjectionCase{259, 300}}) { + constexpr int64_t kRows = 4; + const int64_t width = test_case.width; + const int64_t input_width = test_case.input_width; + std::vector x(kRows * width); + std::vector y(kRows * width); + std::vector act(kRows * input_width); + std::vector weight(input_width); + std::vector expected(kRows * width); + for (int64_t k = 0; k < input_width; ++k) { + weight[k] = BFloat16(static_cast((k % 5) - 2) * 0.25f); + } + for (int64_t r = 0; r < kRows; ++r) { + float dot = 0.0f; for (int64_t k = 0; k < input_width; ++k) { - weight[k] = BFloat16(static_cast((k % 5) - 2) * 0.25f); + act[r * input_width + k] = BFloat16(static_cast(((r + k) % 7) - 3) * 0.125f); + dot = std::fma(act[r * input_width + k].ToFloat(), weight[k].ToFloat(), dot); } - for (int64_t r = 0; r < kRows; ++r) { - float dot = 0.0f; - for (int64_t k = 0; k < input_width; ++k) { - act[r * input_width + k] = BFloat16(static_cast(((r + k) % 7) - 3) * 0.125f); - dot = std::fma(act[r * input_width + k].ToFloat(), weight[k].ToFloat(), dot); - } - const float gate = BFloat16(dot).ToFloat(); - const float exp_value = std::exp(gate >= 0.0f ? -gate : gate); - const float sigmoid = gate >= 0.0f ? 1.0f / (1.0f + exp_value) : exp_value / (1.0f + exp_value); - const BFloat16 denominator(1.0f + BFloat16(exp_value).ToFloat()); - const BFloat16 reciprocal(1.0f / denominator.ToFloat()); - const BFloat16 rounded_sigmoid = - sigmoid_in_fp32 ? BFloat16(sigmoid) - : (gate > 0.0f ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat())); - for (int64_t c = 0; c < width; ++c) { - const int64_t idx = r * width + c; - x[idx] = BFloat16(static_cast((idx % 17) - 8) * 0.13f); - y[idx] = BFloat16(static_cast((idx % 11) - 5) * 0.37f); - const BFloat16 product(y[idx].ToFloat() * rounded_sigmoid.ToFloat()); - expected[idx] = BFloat16(x[idx].ToFloat() + product.ToFloat()); - } + const float gate = BFloat16(dot).ToFloat(); + const float exp_value = std::exp(gate >= 0.0f ? -gate : gate); + const float sigmoid = gate >= 0.0f ? 1.0f / (1.0f + exp_value) : exp_value / (1.0f + exp_value); + const BFloat16 rounded_sigmoid(sigmoid); + for (int64_t c = 0; c < width; ++c) { + const int64_t idx = r * width + c; + x[idx] = BFloat16(static_cast((idx % 17) - 8) * 0.13f); + y[idx] = BFloat16(static_cast((idx % 11) - 5) * 0.37f); + const BFloat16 product(y[idx].ToFloat() * rounded_sigmoid.ToFloat()); + expected[idx] = BFloat16(x[idx].ToFloat() + product.ToFloat()); } - std::vector> providers; - if (auto cuda_ep = DefaultCudaExecutionProvider()) { - providers.push_back(std::move(cuda_ep)); - } - if (providers.empty()) GTEST_SKIP(); - OpTester tester("GatedAdd", 1, onnxruntime::kMSDomain); - tester.AddAttribute("sigmoid_in_fp32", sigmoid_in_fp32 ? 1 : 0); - tester.AddInput("X", {1, kRows, width}, x); - tester.AddInput("Y", {1, kRows, width}, y); - tester.AddInput("gate", {1, kRows, input_width}, act); - tester.AddInput("gate_weight", {input_width, 1}, weight); - tester.AddOutput("output", {1, kRows, width}, expected); - tester.SetOutputTolerance(0.0f, 0.0f); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } + std::vector> providers; + if (auto cuda_ep = DefaultCudaExecutionProvider()) { + providers.push_back(std::move(cuda_ep)); + } + if (providers.empty()) GTEST_SKIP(); + OpTester tester("GatedAdd", 1, onnxruntime::kMSDomain); + tester.AddInput("X", {1, kRows, width}, x); + tester.AddInput("Y", {1, kRows, width}, y); + tester.AddInput("gate", {1, kRows, input_width}, act); + tester.AddInput("gate_weight", {input_width, 1}, weight); + tester.AddOutput("output", {1, kRows, width}, expected); + tester.SetOutputTolerance(0.0f, 0.0f); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } } From ab9f0bcd5ee5fab5c52fc7712ef20e30f1e4b83c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:48:13 +0000 Subject: [PATCH 18/19] Remove dead BFloat16 ArgMax paths and residual null checks; add GatedAdd coverage - Revert BFloat16 hunks in ReduceKernelShared (only instantiated for float/double/MLFloat16 with NO_INDICES, so they were unreachable). The ReduceComputeCore BFloat16 ArgMax/ArgMin cast path stays: it is reached by the non-last-axis fallback. - Drop residual == nullptr branches in ProjectedSigmoidGateMulKernel (X is a required input). - Add tests for fp16/bf16 GatedAdd with hidden size not a multiple of 8 and projected gate with K=2560. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/gated_add_impl.cu | 28 ++++++------------- .../providers/cuda/reduction/reduction_ops.cc | 7 ++--- .../test/contrib_ops/gated_add_op_test.cc | 14 +++++++++- 3 files changed, 25 insertions(+), 24 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu index 2ce68dd290416..82aa1baf8d463 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu @@ -62,29 +62,22 @@ __global__ void ProjectedSigmoidGateMulKernel( if ((output_width & 7) == 0 && (reinterpret_cast(attention + row_base) & 15) == 0 && (reinterpret_cast(output + row_base) & 15) == 0 && - (residual == nullptr || (reinterpret_cast(residual + row_base) & 15) == 0)) { + (reinterpret_cast(residual + row_base) & 15) == 0) { const int64_t vec_width = output_width >> 3; const uint4* attn_vec = reinterpret_cast(attention + row_base); - const uint4* res_vec = residual != nullptr ? reinterpret_cast(residual + row_base) : nullptr; + const uint4* res_vec = reinterpret_cast(residual + row_base); uint4* out_vec = reinterpret_cast(output + row_base); for (int64_t v = threadIdx.x; v < vec_width; v += kProjectedGateThreads) { const uint4 a_u4 = attn_vec[v]; + const uint4 r_u4 = res_vec[v]; const __nv_bfloat16* a_bf = reinterpret_cast(&a_u4); + const __nv_bfloat16* r_bf = reinterpret_cast(&r_u4); uint4 o_u4; __nv_bfloat16* o_bf = reinterpret_cast<__nv_bfloat16*>(&o_u4); - if (res_vec != nullptr) { - const uint4 r_u4 = res_vec[v]; - const __nv_bfloat16* r_bf = reinterpret_cast(&r_u4); #pragma unroll - for (int m = 0; m < 8; ++m) { - const __nv_bfloat16 scaled = __float2bfloat16_rn(__bfloat162float(a_bf[m]) * gate); - o_bf[m] = __float2bfloat16_rn(__bfloat162float(scaled) + __bfloat162float(r_bf[m])); - } - } else { -#pragma unroll - for (int m = 0; m < 8; ++m) { - o_bf[m] = __float2bfloat16_rn(__bfloat162float(a_bf[m]) * gate); - } + for (int m = 0; m < 8; ++m) { + const __nv_bfloat16 scaled = __float2bfloat16_rn(__bfloat162float(a_bf[m]) * gate); + o_bf[m] = __float2bfloat16_rn(__bfloat162float(scaled) + __bfloat162float(r_bf[m])); } out_vec[v] = o_u4; } @@ -92,11 +85,8 @@ __global__ void ProjectedSigmoidGateMulKernel( } for (int64_t column = threadIdx.x; column < output_width; column += kProjectedGateThreads) { const int64_t index = row_base + column; - __nv_bfloat16 value = __float2bfloat16_rn(__bfloat162float(attention[index]) * gate); - if (residual != nullptr) { - value = __float2bfloat16_rn(__bfloat162float(value) + __bfloat162float(residual[index])); - } - output[index] = value; + const __nv_bfloat16 scaled = __float2bfloat16_rn(__bfloat162float(attention[index]) * gate); + output[index] = __float2bfloat16_rn(__bfloat162float(scaled) + __bfloat162float(residual[index])); } } diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index a110cb1e55ce2..18cc5fdee0668 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -97,9 +97,8 @@ Status ReduceKernel::ReduceKernelShared( int64_t input_count = input_shape.Size(); IAllocatorUniquePtr temp_X; - if (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES && - (std::is_same::value || std::is_same::value)) { - // ArgMax/ArgMin with FP16/BF16 are not supported by cudnn, so convert input to fp32 then call cudnn + if (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES && std::is_same::value) { + // ArgMax/ArgMin with FP16 are not supported by cudnn, so convert input to fp32 then call cudnn temp_X = GetScratchBuffer(input_count, stream); cudnn_type_X = CUDNN_DATA_FLOAT; Impl_Cast(cuda_stream, reinterpret_cast(X), temp_X.get(), input_shape.Size()); @@ -115,7 +114,7 @@ Status ReduceKernel::ReduceKernelShared( } CudnnReduceDescriptor reduce_desc; - if constexpr (std::is_same::value || std::is_same::value) + if constexpr (std::is_same::value) ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, CudnnTensor::GetDataType(), ReduceTensorIndices)); else ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, cudnn_type_X, ReduceTensorIndices)); diff --git a/onnxruntime/test/contrib_ops/gated_add_op_test.cc b/onnxruntime/test/contrib_ops/gated_add_op_test.cc index 24407becd066b..d0ca2591b19a1 100644 --- a/onnxruntime/test/contrib_ops/gated_add_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_add_op_test.cc @@ -162,6 +162,17 @@ TEST(ContribOpGatedAddTest, BFloat16) { RunGatedAddTest({1, 4, 2048}); } +TEST(ContribOpGatedAddTest, Float16NonVectorizedHiddenSize) { + RunGatedAddTest({3, 5, 7}); +} + +TEST(ContribOpGatedAddTest, BFloat16NonVectorizedHiddenSize) { + if (!CudaHasBF16Support()) { + GTEST_SKIP() << "bfloat16 requires compute capability 8.0 or later"; + } + RunGatedAddTest({3, 5, 7}); +} + TEST(ContribOpGatedAddTest, MismatchedYShape) { auto execution_providers = AvailableGatedOpExecutionProviders(); for (auto& ep : execution_providers) { @@ -187,7 +198,8 @@ TEST(ContribOpGatedAddTest, ProjectedGateBFloat16) { int64_t width; int64_t input_width; }; - for (const ProjectionCase& test_case : {ProjectionCase{8, 16}, ProjectionCase{2560, 16}, ProjectionCase{259, 300}}) { + for (const ProjectionCase& test_case : {ProjectionCase{8, 16}, ProjectionCase{2560, 16}, ProjectionCase{2560, 2560}, + ProjectionCase{259, 300}}) { constexpr int64_t kRows = 4; const int64_t width = test_case.width; const int64_t input_width = test_case.input_width; From edf701c01c15548c99a0881dd24fa72eb5f5e441 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:48:51 +0000 Subject: [PATCH 19/19] Document plain and projected modes of GatedAdd Rewrite the GatedAdd schema doc and input descriptions to cover both the direct-gate mode and the projected-gate mode (gate_weight), including the FP32 accumulation, sigmoid and BFloat16 rounding points, shape contract and CUDA + BFloat16 support. Update docs/ContribOperators.md to match. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 33 +++++++++++--- .../core/graph/contrib_ops/bert_defs.cc | 43 +++++++++++++++---- 2 files changed, 61 insertions(+), 15 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 1ba3d71924fab..1ebe6d2e2f250 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2387,13 +2387,32 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.GatedAdd** - Adds one tensor to another tensor scaled by a per-row gate: + Adds one tensor to another tensor scaled by a per-row gate. X and Y have shape (..., C). The + leading dimensions of gate must match the leading dimensions of X (there is no broadcasting over + them); the optional gate_weight input selects how the per-row gate is obtained. + + Plain mode (gate_weight is not provided): gate has shape (..., 1) and is used directly as the + gate, broadcast over C. No sigmoid is applied. output = X + round_to_T(Y * gate) - X and Y have shape (..., C), and gate has shape (..., 1). The gate is broadcast - over C. For reduced-precision types, the product is rounded to T before the add, - matching separate ONNX Mul and Add operators. + For reduced-precision types, the product is rounded to T before the add, matching separate ONNX + Mul and Add operators. This mode supports float, float16 and bfloat16, subject to execution provider kernel support. + + Projected mode (gate_weight is provided): gate holds an activation with shape (..., K) and + gate_weight has shape (K, 1). For every row, the gate is computed by projecting the activation + with gate_weight and applying a sigmoid: + + p = round_to_bfloat16(sum_k(gate[..., k] * gate_weight[k])) + s = round_to_bfloat16(sigmoid(p)) + output = round_to_bfloat16(X + round_to_bfloat16(Y * s)) + + The dot product, the sigmoid and the multiplication and addition are evaluated in FP32 (the + sigmoid is always computed in FP32, and the order of the FP32 accumulation over K is not + specified); the result is rounded to bfloat16 after the projection, after the sigmoid, after the + multiplication and after the addition. This mode is only implemented by the CUDA execution + provider for bfloat16 tensors. Other execution providers and types return NOT_IMPLEMENTED when + gate_weight is provided. #### Version @@ -2405,11 +2424,11 @@ This version of the operator has been available since version 1 of the 'com.micr
X : T
Unscaled input with shape (..., C).
Y : T
-
Input scaled by gate, with the same shape as X.
+
Input scaled by the gate, with the same shape as X.
gate : T
-
Per-row gate with shape (..., 1), or activation (..., K) when gate_weight is provided.
+
In plain mode, the per-row gate with shape (..., 1). When gate_weight is provided, the activation with shape (..., K) that is projected with gate_weight and passed through a sigmoid to produce the per-row gate. Its leading dimensions must match X.
gate_weight (optional) : T
-
Optional projection weight [K, 1] applied to gate before sigmoid.
+
Optional projection weight with shape (K, 1), where K is the last dimension of gate. When present, the gate is sigmoid(gate . gate_weight) computed per row in FP32 (projected mode). Only supported by the CUDA execution provider with bfloat16.
#### Outputs diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index aea6928e8c3a5..aca69d9276277 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -6105,23 +6105,50 @@ ONNX_MS_OPERATOR_SET_SCHEMA( })); constexpr const char* GatedAdd_ver1_doc = R"DOC( -Adds one tensor to another tensor scaled by a per-row gate: +Adds one tensor to another tensor scaled by a per-row gate. X and Y have shape (..., C). The +leading dimensions of gate must match the leading dimensions of X (there is no broadcasting over +them); the optional gate_weight input selects how the per-row gate is obtained. + +Plain mode (gate_weight is not provided): gate has shape (..., 1) and is used directly as the +gate, broadcast over C. No sigmoid is applied. output = X + round_to_T(Y * gate) -X and Y have shape (..., C), and gate has shape (..., 1). The gate is broadcast -over C. For reduced-precision types, the product is rounded to T before the add, -matching separate ONNX Mul and Add operators. +For reduced-precision types, the product is rounded to T before the add, matching separate ONNX +Mul and Add operators. This mode supports float, float16 and bfloat16, subject to execution provider kernel support. + +Projected mode (gate_weight is provided): gate holds an activation with shape (..., K) and +gate_weight has shape (K, 1). For every row, the gate is computed by projecting the activation +with gate_weight and applying a sigmoid: + + p = round_to_bfloat16(sum_k(gate[..., k] * gate_weight[k])) + s = round_to_bfloat16(sigmoid(p)) + output = round_to_bfloat16(X + round_to_bfloat16(Y * s)) + +The dot product, the sigmoid and the multiplication and addition are evaluated in FP32 (the +sigmoid is always computed in FP32, and the order of the FP32 accumulation over K is not +specified); the result is rounded to bfloat16 after the projection, after the sigmoid, after the +multiplication and after the addition. This mode is only implemented by the CUDA execution +provider for bfloat16 tensors. Other execution providers and types return NOT_IMPLEMENTED when +gate_weight is provided. )DOC"; ONNX_MS_OPERATOR_SET_SCHEMA( GatedAdd, 1, OpSchema() .SetDoc(GatedAdd_ver1_doc) - .Input(0, "X", "Unscaled input with shape (..., C).", "T") - .Input(1, "Y", "Input scaled by gate, with the same shape as X.", "T") - .Input(2, "gate", "Per-row gate with shape (..., 1), or activation (..., K) when gate_weight is provided.", "T") - .Input(3, "gate_weight", "Optional projection weight [K, 1] applied to gate before sigmoid.", "T", OpSchema::Optional) + .Input(0, "X", + "Unscaled input with shape (..., C).", + "T") + .Input(1, "Y", + "Input scaled by the gate, with the same shape as X.", + "T") + .Input(2, "gate", + "In plain mode, the per-row gate with shape (..., 1). When gate_weight is provided, the activation with shape (..., K) that is projected with gate_weight and passed through a sigmoid to produce the per-row gate. Its leading dimensions must match X.", + "T") + .Input(3, "gate_weight", + "Optional projection weight with shape (K, 1), where K is the last dimension of gate. When present, the gate is sigmoid(gate . gate_weight) computed per row in FP32 (projected mode). Only supported by the CUDA execution provider with bfloat16.", + "T", OpSchema::Optional) .Output(0, "output", "Gated sum with the same shape as X.", "T") .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"},