Add more kernels for Qwen-3.5 ops - #32106
Merged
Merged
Conversation
…nGate Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Copilot created this pull request from a session on behalf of
kunal-vaishnavi
August 15, 2026 04:39
View session
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Adds CPU and WebGPU execution support for Qwen-3.5 gated attention operators.
Changes:
- Implements GatedAdd, GatedRMSNorm, and LinearAttentionGate kernels.
- Registers the new CPU and WebGPU kernels.
- Extends shared tests across CPU, CUDA, and WebGPU.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
onnxruntime/test/contrib_ops/linear_attention_gates_op_test.cc |
Runs gate and RMSNorm tests across supported EPs. |
onnxruntime/test/contrib_ops/gated_add_op_test.cc |
Extends GatedAdd tests to CPU and WebGPU. |
onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc |
Registers WebGPU kernels. |
onnxruntime/contrib_ops/webgpu/bert/linear_attention_gates.h |
Declares WebGPU gate programs and kernels. |
onnxruntime/contrib_ops/webgpu/bert/linear_attention_gates.cc |
Implements WebGPU gate and RMSNorm shaders. |
onnxruntime/contrib_ops/webgpu/bert/gated_add.h |
Declares the WebGPU GatedAdd kernel. |
onnxruntime/contrib_ops/webgpu/bert/gated_add.cc |
Implements the WebGPU GatedAdd shader. |
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc |
Registers CPU kernels. |
onnxruntime/contrib_ops/cpu/bert/linear_attention_gates.h |
Declares CPU gate kernels. |
onnxruntime/contrib_ops/cpu/bert/linear_attention_gates.cc |
Implements CPU gate and RMSNorm kernels. |
onnxruntime/contrib_ops/cpu/bert/gated_add.h |
Declares the CPU GatedAdd kernel. |
onnxruntime/contrib_ops/cpu/bert/gated_add.cc |
Implements the CPU GatedAdd kernel. |
Suppressed comments (1)
onnxruntime/contrib_ops/cpu/bert/linear_attention_gates.cc:157
- Narrowing
num_rowstoint32_tis unsafe becauseTryBatchParallelForsupportsstd::ptrdiff_t. With more thanINT32_MAXnormalization groups, the cast can produce a non-positive total, so no work runs and the uninitialized output is returned as successful. Preserve the wider count.
context->GetOperatorThreadPool(), static_cast<int32_t>(num_rows),
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Tianlei Wu (tianleiwu)
requested changes
Aug 15, 2026
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
Contributor
There was a problem hiding this comment.
Two deterministic, PR-caused CI failures remain: the Float16 GatedAdd test overrides the repository's reduced-precision tolerance, and the generated CPU kernel documentation is stale. Details are inline. The implementation review otherwise found the operator validation, QMoE checked arithmetic and tiling, CUDA registrations, and targeted coverage well structured.
kunal-vaishnavi
force-pushed
the
copilot/add-cuda-ops-for-qwen-3-5
branch
from
August 15, 2026 21:08
abb5b8d to
38f8596
Compare
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…ache hints, unify GatedAdd path Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Tianlei Wu (tianleiwu)
added a commit
to microsoft/onnxruntime-genai
that referenced
this pull request
Aug 21, 2026
## Description Fuse Qwen3.6 MoE shared-expert scaling and routed/shared addition into the `com.microsoft::GatedAdd` contrib operator. Qwen3.6 shared-expert Mul+Add is replaced by com.microsoft::GatedAdd for CPU, CUDA, and WebGPU. Unsupported execution providers retain portable ONNX `Mul` + `Add`. No `fuse_shared_expert_gate` option remains. This PR is stacked on #2353, which is itself stacked on #2351. Dependencies are microsoft/onnxruntime#31835 for CUDA and merged microsoft/onnxruntime#32106 for CPU/WebGPU. ## Changes - Return the shared-expert projection and scalar gate separately from `make_shared_expert`. - Emit one `GatedAdd` per MoE layer on CPU, CUDA, and WebGPU. - Preserve the standard `Mul` + `Add` graph as an explicit fallback for unsupported EPs. - Add focused tests for fused and fallback graph construction. ## Performance On Qwen3.6-35B-A3B-NVFP4 with N=3 MTP, the real exported graph replaces 40 main-model pairs plus one MTP pair. Counterbalanced H200 measurements reduced median decode latency from 7.311 to 7.225 ms/round (-1.18%). Graph-off Nsight measured 40.35 fewer launches/round and 1.30% lower GPU kernel time. ## Validation - `50 passed`: `test_precision.py` plus `test_qwen_gated_add.py`. - Real graph census: 40 `GatedAdd` nodes in `text.onnx`, one in `mtp.onnx`. - After normalizing fused edge names, all other nodes, initializers, graph inputs, and graph outputs are unchanged from the baseline export. - Runtime float, FP16, and BF16 results are bit-exact with separate `Mul` + `Add`. ## Stack note The source commit `5c35bb02fc` also contained an unrelated MTP prefill chunk default change in `src/mtp_generator.{h,cpp}`. Those files are intentionally excluded here because they depend on runtime PR #2352 rather than builder PR #2353.
This was referenced Sep 10, 2026
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Description
Adds Qwen-3.5-related contrib kernel coverage across CPU, WebGPU, and CUDA. The PR introduces CPU and WebGPU implementations for
GatedAdd,LinearAttentionGate, andGatedRMSNorm, and enables CUDABFloat16registrations forCausalConvWithState,LinearAttention, andLpNormalization.Changes
GatedAdd,LinearAttentionGate, andGatedRMSNormcontrib kernels and registered them in the CPU contrib kernel table.GatedAdd,LinearAttentionGate, andGatedRMSNormcontrib kernels, registered them in the WebGPU contrib kernel table, and fixedLinearAttentionGateshader binding order.BFloat16kernel registrations forCausalConvWithStateandLinearAttention, plusLpNormalizationBF16 support starting at opset 22.GatedAdd,LinearAttentionGate,GatedRMSNorm,CausalConvWithState, andLinearAttention, and added a CUDA capability-gated BF16LpNormalizationtest.Motivation and Context
These changes add kernel and type coverage needed by Qwen-3.5-style model patterns across supported execution providers, including BF16 CUDA paths for the affected ops.