-
Notifications
You must be signed in to change notification settings - Fork 4.2k
Add more kernels for Qwen-3.5 ops #32106
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
20 commits
Select commit
Hold shift + click to select a range
adddd06
Add CPU and WebGPU kernels for GatedAdd, GatedRMSNorm, LinearAttentio…
Copilot c35a2a0
Fix WebGPU input/output binding order and remove unused include
Copilot d336d6e
Reuse shared activation utilities for Qwen CPU gates
Copilot 72d1767
Document shared activation functor reuse
Copilot a2302a0
Align LinearAttentionGate CPU error message
Copilot 8661dc5
Align WebGPU LinearAttentionGate error message
Copilot ed1fc02
Apply requested clang-format corrections
Copilot e8e303f
Clarify WebGPU SiLU helper name
Copilot e6216fd
Fix CPU iteration width and FP16 rounding
Copilot d00553c
Clarify WebGPU sigmoid helper
Copilot 52258f0
Cast WebGPU FP16 GatedAdd output
Copilot 72e5dc6
Fix float buffer packing into Eigen for CPU
Copilot d7c787f
Fix clang-format indentation in linear_attention_gates
Copilot cdba354
Add CUDA BF16 LpNormalization support and test
Copilot 2539fa3
Add missing registration for LpNorm BF16 CUDA
Copilot 50e329e
Add CUDA BF16 support for CausalConvWithState and LinearAttention
Copilot 38f8596
Update docs
kunal-vaishnavi fc78255
Relax GatedAdd reduced precision tolerance
Copilot b96bf7f
Relax WebGPU GatedAdd FP16 tolerance
Copilot 1460571
Address WebGPU review feedback: use Get/SetByOffset, drop redundant c…
Copilot File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,95 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
|
|
||
| #include "contrib_ops/cpu/bert/gated_add.h" | ||
|
|
||
| #include "core/framework/tensor.h" | ||
| #include "core/platform/threadpool.h" | ||
|
|
||
| namespace onnxruntime { | ||
| namespace contrib { | ||
|
|
||
| #define REGISTER_KERNEL_TYPED(T) \ | ||
| ONNX_OPERATOR_TYPED_KERNEL_EX( \ | ||
| GatedAdd, \ | ||
| kMSDomain, \ | ||
| 1, \ | ||
| T, \ | ||
| kCpuExecutionProvider, \ | ||
| KernelDefBuilder() \ | ||
| .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \ | ||
| GatedAdd<T>); | ||
|
|
||
| REGISTER_KERNEL_TYPED(float) | ||
| REGISTER_KERNEL_TYPED(MLFloat16) | ||
|
|
||
| #undef REGISTER_KERNEL_TYPED | ||
|
|
||
| namespace { | ||
|
|
||
| // output = x + round_to_T(y * gate). For MLFloat16 the product is rounded to half before the | ||
| // add, matching separate ONNX Mul and Add operators. | ||
| template <typename T> | ||
| inline T GatedAddValue(T x, T y, T gate) { | ||
| if constexpr (std::is_same_v<T, MLFloat16>) { | ||
| const T product(y.ToFloat() * gate.ToFloat()); | ||
| return T(x.ToFloat() + product.ToFloat()); | ||
| } else { | ||
| return x + y * gate; | ||
| } | ||
| } | ||
|
|
||
| } // namespace | ||
|
|
||
| template <typename T> | ||
| Status GatedAdd<T>::Compute(OpKernelContext* context) const { | ||
| const Tensor* x = context->Input<Tensor>(0); | ||
| const Tensor* y = context->Input<Tensor>(1); | ||
| const Tensor* gate = context->Input<Tensor>(2); | ||
| const TensorShape& shape = x->Shape(); | ||
|
|
||
| ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1, "X must have rank >= 1"); | ||
| ORT_RETURN_IF_NOT(y->Shape() == shape, "Y must have the same shape as X"); | ||
| ORT_RETURN_IF_NOT(gate->Shape().NumDimensions() == shape.NumDimensions(), | ||
| "gate must have the same rank as X"); | ||
|
|
||
| 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); | ||
| const int64_t count = shape.Size(); | ||
| if (count == 0) { | ||
| return Status::OK(); | ||
| } | ||
|
|
||
| const T* x_data = x->Data<T>(); | ||
| const T* y_data = y->Data<T>(); | ||
| const T* gate_data = gate->Data<T>(); | ||
| T* output_data = output->MutableData<T>(); | ||
| const int64_t num_rows = count / hidden_size; | ||
|
|
||
| concurrency::ThreadPool::TryBatchParallelFor( | ||
| context->GetOperatorThreadPool(), onnxruntime::narrow<ptrdiff_t>(num_rows), | ||
| [&](ptrdiff_t row) { | ||
| const int64_t offset = row * hidden_size; | ||
| const T gate_value = gate_data[row]; | ||
| for (int64_t i = 0; i < hidden_size; ++i) { | ||
| output_data[offset + i] = GatedAddValue<T>(x_data[offset + i], y_data[offset + i], gate_value); | ||
| } | ||
| }, | ||
| 0); | ||
|
|
||
| return Status::OK(); | ||
| } | ||
|
|
||
| template class GatedAdd<float>; | ||
| template class GatedAdd<MLFloat16>; | ||
|
|
||
| } // namespace contrib | ||
| } // namespace onnxruntime | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,21 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
|
|
||
| #pragma once | ||
|
|
||
| #include "core/common/common.h" | ||
| #include "core/framework/op_kernel.h" | ||
|
|
||
| namespace onnxruntime { | ||
| namespace contrib { | ||
|
|
||
| // output = X + round_to_T(Y * gate), with gate broadcast across the last dimension. | ||
| template <typename T> | ||
| class GatedAdd final : public OpKernel { | ||
| public: | ||
| explicit GatedAdd(const OpKernelInfo& info) : OpKernel(info) {} | ||
| Status Compute(OpKernelContext* context) const override; | ||
| }; | ||
|
|
||
| } // namespace contrib | ||
| } // namespace onnxruntime |
181 changes: 181 additions & 0 deletions
181
onnxruntime/contrib_ops/cpu/bert/linear_attention_gates.cc
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,181 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
|
github-advanced-security[bot] marked this conversation as resolved.
Fixed
|
||
| // Licensed under the MIT License. | ||
|
|
||
| #include "contrib_ops/cpu/bert/linear_attention_gates.h" | ||
|
|
||
| #include <cmath> | ||
|
|
||
| #include "core/framework/tensor.h" | ||
| #include "core/mlas/inc/mlas.h" | ||
| #include "core/platform/threadpool.h" | ||
|
|
||
| namespace onnxruntime { | ||
| namespace contrib { | ||
|
|
||
| #define REGISTER_KERNEL_TYPED(Op, T) \ | ||
| ONNX_OPERATOR_TYPED_KERNEL_EX( \ | ||
| Op, \ | ||
| kMSDomain, \ | ||
| 1, \ | ||
| T, \ | ||
| kCpuExecutionProvider, \ | ||
| KernelDefBuilder() \ | ||
| .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \ | ||
| .TypeConstraint("TF", DataTypeImpl::GetTensorType<float>()), \ | ||
| Op<T>); | ||
|
|
||
| REGISTER_KERNEL_TYPED(LinearAttentionGate, float) | ||
| REGISTER_KERNEL_TYPED(LinearAttentionGate, MLFloat16) | ||
|
|
||
| #undef REGISTER_KERNEL_TYPED | ||
|
|
||
| #define REGISTER_KERNEL_TYPED(Op, T) \ | ||
| ONNX_OPERATOR_TYPED_KERNEL_EX( \ | ||
| Op, \ | ||
| kMSDomain, \ | ||
| 1, \ | ||
| T, \ | ||
| kCpuExecutionProvider, \ | ||
| KernelDefBuilder() \ | ||
| .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \ | ||
| Op<T>); | ||
|
|
||
| REGISTER_KERNEL_TYPED(GatedRMSNorm, float) | ||
| REGISTER_KERNEL_TYPED(GatedRMSNorm, MLFloat16) | ||
|
|
||
| #undef REGISTER_KERNEL_TYPED | ||
|
|
||
| namespace { | ||
|
|
||
| inline float SigmoidFloat(float value) { | ||
| float output; | ||
| MlasComputeLogistic(&value, &output, 1); | ||
| return output; | ||
| } | ||
|
|
||
| inline float SoftplusFloat(float value) { | ||
| return value > 0.0f ? value + std::log(std::exp(-value) + 1.0f) : std::log(std::exp(value) + 1.0f); | ||
| } | ||
|
|
||
| } // namespace | ||
|
|
||
| template <typename T> | ||
| Status LinearAttentionGate<T>::Compute(OpKernelContext* context) const { | ||
| const Tensor* a = context->Input<Tensor>(0); | ||
| const Tensor* dt_bias = context->Input<Tensor>(1); | ||
| const Tensor* decay_scale = context->Input<Tensor>(2); | ||
| const Tensor* b = context->Input<Tensor>(3); // optional | ||
|
|
||
| const auto& a_shape = a->Shape(); | ||
| ORT_RETURN_IF_NOT(a_shape.NumDimensions() >= 1, "a must have rank >= 1"); | ||
| const int64_t num_heads = a_shape[a_shape.NumDimensions() - 1]; | ||
| ORT_RETURN_IF_NOT(num_heads > 0, "a last dimension must be positive"); | ||
|
|
||
| ORT_RETURN_IF_NOT(dt_bias->Shape().Size() == num_heads, | ||
| "dt_bias must have ", num_heads, " elements, got ", dt_bias->Shape().Size()); | ||
| ORT_RETURN_IF_NOT(decay_scale->Shape().Size() == num_heads, | ||
| "decay_scale must have ", num_heads, " elements, got ", decay_scale->Shape().Size()); | ||
|
|
||
| Tensor* decay = context->Output(0, a_shape); | ||
| Tensor* beta = context->Output(1, a_shape); | ||
|
|
||
| if (beta != nullptr) { | ||
| ORT_RETURN_IF_NOT(b != nullptr, "The b input is required when the beta output is requested"); | ||
| ORT_RETURN_IF_NOT(b->Shape() == a_shape, "b must have the same shape as a"); | ||
| } | ||
|
|
||
| const int64_t count = a_shape.Size(); | ||
| if (count == 0) { | ||
| return Status::OK(); | ||
| } | ||
|
|
||
| const T* a_data = a->Data<T>(); | ||
| const T* b_data = b == nullptr ? nullptr : b->Data<T>(); | ||
| const float* dt_bias_data = dt_bias->Data<float>(); | ||
| const float* decay_scale_data = decay_scale->Data<float>(); | ||
| T* decay_data = decay->MutableData<T>(); | ||
| T* beta_data = beta == nullptr ? nullptr : beta->MutableData<T>(); | ||
|
|
||
| const int64_t num_tokens = count / num_heads; | ||
|
|
||
| concurrency::ThreadPool::TryBatchParallelFor( | ||
| context->GetOperatorThreadPool(), onnxruntime::narrow<ptrdiff_t>(num_tokens), | ||
| [&](ptrdiff_t token) { | ||
| const int64_t offset = token * num_heads; | ||
| for (int64_t h = 0; h < num_heads; ++h) { | ||
| const int64_t idx = offset + h; | ||
| const float biased = static_cast<float>(a_data[idx]) + dt_bias_data[h]; | ||
| decay_data[idx] = static_cast<T>(decay_scale_data[h] * SoftplusFloat(biased)); | ||
| if (beta_data != nullptr) { | ||
| beta_data[idx] = static_cast<T>(SigmoidFloat(static_cast<float>(b_data[idx]))); | ||
| } | ||
| } | ||
| }, | ||
| 0); | ||
|
|
||
| return Status::OK(); | ||
| } | ||
|
|
||
| template <typename T> | ||
| GatedRMSNorm<T>::GatedRMSNorm(const OpKernelInfo& info) : OpKernel(info) { | ||
| epsilon_ = info.GetAttrOrDefault<float>("epsilon", 1e-5f); | ||
| } | ||
|
|
||
| template <typename T> | ||
| Status GatedRMSNorm<T>::Compute(OpKernelContext* context) const { | ||
| const Tensor* input = context->Input<Tensor>(0); | ||
| const Tensor* scale = context->Input<Tensor>(1); | ||
| const Tensor* gate = context->Input<Tensor>(2); | ||
|
|
||
| const auto& shape = input->Shape(); | ||
| ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1, "X must have rank >= 1"); | ||
| ORT_RETURN_IF_NOT(gate->Shape() == shape, "gate must have the same shape as X"); | ||
|
|
||
| const int64_t norm_size = scale->Shape().Size(); | ||
| ORT_RETURN_IF_NOT(norm_size > 0, "scale must not be empty"); | ||
| const int64_t last_dim = shape[shape.NumDimensions() - 1]; | ||
| ORT_RETURN_IF_NOT(last_dim % norm_size == 0, | ||
| "X last dimension (", last_dim, ") must be a multiple of the scale length (", | ||
| norm_size, ")"); | ||
|
|
||
| Tensor* output = context->Output(0, shape); | ||
| const int64_t count = shape.Size(); | ||
| if (count == 0) { | ||
| return Status::OK(); | ||
| } | ||
| const int64_t num_rows = count / norm_size; | ||
|
|
||
| const T* input_data = input->Data<T>(); | ||
| const T* scale_data = scale->Data<T>(); | ||
| const T* gate_data = gate->Data<T>(); | ||
| T* output_data = output->MutableData<T>(); | ||
|
|
||
| concurrency::ThreadPool::TryBatchParallelFor( | ||
| context->GetOperatorThreadPool(), onnxruntime::narrow<ptrdiff_t>(num_rows), | ||
| [&](ptrdiff_t row) { | ||
| const int64_t offset = row * norm_size; | ||
| float sum_sq = 0.0f; | ||
| for (int64_t i = 0; i < norm_size; ++i) { | ||
| const float v = static_cast<float>(input_data[offset + i]); | ||
| sum_sq += v * v; | ||
| } | ||
| const float inv_rms = 1.0f / std::sqrt(sum_sq / static_cast<float>(norm_size) + epsilon_); | ||
| for (int64_t i = 0; i < norm_size; ++i) { | ||
| const float z = static_cast<float>(gate_data[offset + i]); | ||
| const float normalized = static_cast<float>(input_data[offset + i]) * inv_rms * | ||
| static_cast<float>(scale_data[i]); | ||
| output_data[offset + i] = static_cast<T>(normalized * (z * SigmoidFloat(z))); | ||
| } | ||
| }, | ||
| 0); | ||
|
|
||
| return Status::OK(); | ||
| } | ||
|
|
||
| template class LinearAttentionGate<float>; | ||
| template class LinearAttentionGate<MLFloat16>; | ||
| template class GatedRMSNorm<float>; | ||
| template class GatedRMSNorm<MLFloat16>; | ||
|
|
||
| } // namespace contrib | ||
| } // namespace onnxruntime | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.