[WebGPU] QKV and MLP layer fusions for Qwen3-style models

[hariharans29]

Description

Summary

Adds two WebGPU-only graph fusions and the contrib ops they target, plus a small
refactor of the existing MatMulNBits dispatch logic so the new fused kernels
can share its predicates.

Component	Files	Purpose
MatMulNBitsMlp op + kernel	contrib_ops/webgpu/quantization/matmul_nbits_mlp.{cc,h}, *.wgsl.template (3)	Fuses the SwiGLU MLP block: optional (Skip)SimplifiedLayerNormalization + two MatMulNBits projections (gate, up) + optional biases + Sigmoid/Mul (SiLU) + element-wise Mul. Single dispatch instead of 5–7.
MatMulNBitsQkv op + kernel	contrib_ops/webgpu/quantization/matmul_nbits_qkv.{cc,h}, *.wgsl.template	Fuses (Skip)SimplifiedLayerNormalization + three MatMulNBits projections (Q, K, V) sharing the same input. Single dispatch instead of 4.
Op schemas	core/graph/contrib_ops/contrib_defs.cc	MatMulNBitsMlp and MatMulNBitsQkv contrib op schemas (kMSDomain, opset 1).
Graph transformers	core/optimizer/matmul_nbits_{mlp,qkv}_fusion.{cc,h}	Pattern-match the source subgraphs and emit the fused ops. EP-gated to WebGPU only — no impact on other EPs. Registered in graph_transformer_utils.cc.
Dispatch helpers	contrib_ops/webgpu/quantization/matmul_nbits_common.{cc,h} + matmul_nbits.cc	Extracts the "would this dispatch use Subgroup-Matrix / DP4A / WideTile?" predicates into pure functions reusable by the fused kernels. No behavior change in the unfused MatMulNBits path.
Tests	test/optimizer/matmul_nbits_{mlp,qkv}_fusion_test.cc, graph_transform_utils_test.cc	Unit tests for the new transformers (positive + negative cases).

Motivation and Context

~25-30% decode TPS throughput improvement on WebGPU + D3D backend on Windows. GPU used: RTX 5060Ti for Qwe3-1.7B.

BEFORE (95 decode TPS): main branch

AFTER (120+ decode TPS): PR branch

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

This PR adds WebGPU-focused fused operators and optimizer passes for decoder-style MatMulNBits patterns (MLP gate/up and QKV projections), along with tests and a microbenchmark to evaluate decode performance/correctness.

Changes:

• Introduces new contrib ops MatMulNBitsMlp and MatMulNBitsQkv (schemas + WebGPU kernels + WGSL templates).
• Adds graph transformers MatMulNBitsMlpFusion / MatMulNBitsQkvFusion and corresponding optimizer tests.
• Improves WebGPU runtime support (graph-capture buffer manager activation, queue-idle wait helper, better shader compilation diagnostics) and adds a decode microbenchmark.

Reviewed changes

Copilot reviewed 33 out of 33 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
onnxruntime/test/optimizer/matmul_nbits_qkv_fusion_test.cc	New unit tests validating QKV fusion and output contracts on WebGPU.
onnxruntime/test/optimizer/matmul_nbits_mlp_fusion_test.cc	New unit tests validating MLP fusion (simplified/skip + passthrough) on WebGPU.
onnxruntime/test/optimizer/graph_transform_utils_test.cc	Minor formatting-only tweak (blank line).
onnxruntime/test/onnx/microbenchmark/webgpu_matmul_nbits_decode.cc	New benchmark harness for fused/unfused decode paths on WebGPU.
onnxruntime/test/onnx/microbenchmark/main.cc	Adjusts benchmark env logging severity.
onnxruntime/core/session/ort_version_check.h	Makes version parsing consteval-friendly with a macro fallback.
onnxruntime/core/providers/webgpu/webgpu_execution_provider.h	Tracks when graph-capture buffer manager is active.
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc	Lazily creates/activates graph buffer manager for capture; allocator uses dynamic buffer manager getter.
onnxruntime/core/providers/webgpu/webgpu_context.h	Adds WaitForQueueIdle() declaration.
onnxruntime/core/providers/webgpu/webgpu_context.cc	Implements WaitForQueueIdle() using OnSubmittedWorkDone.
onnxruntime/core/providers/webgpu/program_manager.cc	Enhances pipeline build failures with shader compilation diagnostics.
onnxruntime/core/providers/webgpu/compute_context.h	Adds FlushAndWait() convenience for flushing + waiting on queue idle.
onnxruntime/core/providers/webgpu/allocator.h	Adds allocator ctor that accepts a buffer-manager getter function.
onnxruntime/core/providers/webgpu/allocator.cc	Implements getter-based allocator to support switching buffer managers.
onnxruntime/core/optimizer/matmul_nbits_qkv_fusion.h	New transformer declaration for QKV fusion.
onnxruntime/core/optimizer/matmul_nbits_qkv_fusion.cc	New transformer implementation for QKV fusion.
onnxruntime/core/optimizer/matmul_nbits_mlp_fusion.h	New transformer declaration for MLP fusion.
onnxruntime/core/optimizer/matmul_nbits_mlp_fusion.cc	New transformer implementation for MLP fusion.
onnxruntime/core/optimizer/graph_transformer_utils.cc	Registers the new fusion transformers.
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Adds contrib operator schemas/docs for MatMulNBitsMlp and MatMulNBitsQkv.
onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc	Registers WebGPU kernels for the new fused ops.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.wgsl.template	New WGSL template implementing fused QKV decode kernel.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h	New WebGPU kernel wrapper for MatMulNBitsQkv.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc	New WebGPU kernel implementation for MatMulNBitsQkv.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp_wide_tile_m1.wgsl.template	New WGSL template for an MLP wide-tile variant.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.wgsl.template	New WGSL template implementing fused MLP (optionally with norm/skip/passthrough).
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h	New WebGPU kernel wrapper for MatMulNBitsMlp.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc	New WebGPU kernel implementation for MatMulNBitsMlp.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_common.h	Adds declarations for “would apply” dispatch-selection helpers and shared constants.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_common.cc	Implements the new dispatch-selection helpers.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc	Refactors path selection to use the new “would apply” helpers.
onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_mlp.wgsl.template	Adds WGSL template for DP4A MLP path.
cmake/onnxruntime_unittests.cmake	Wires the new WebGPU decode benchmark into the benchmark target sources.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[hariharans29]

Review: [WebGPU] QKV and MLP Layer Fusions for Qwen3-style Models

Summary

This PR adds two WebGPU-only graph fusions (MatMulNBitsMlp and MatMulNBitsQkv) that fuse normalization + quantized matmul projections into single GPU dispatches. It also refactors dispatch decision predicates into shared WouldApply* helpers in matmul_nbits_common.{cc,h}, and extracts RunSkipLayerNormProgram / RunLayerNormProgram so fused ops can reuse existing norm dispatch logic. The claimed result is 25-30% decode TPS improvement.

Overall Assessment: Positive with suggestions

The high-level design is sound — fusing norm + multiple matmuls into single dispatches is a well-established optimization for LLM inference. The code is well-commented, the pattern matching is thorough, and both graph-level and runtime correctness tests are included.

🟢 Strengths

1. Clean separation of concerns: Graph transformers (pattern matching) live in core/optimizer/, kernel implementations in contrib_ops/webgpu/, and shared dispatch logic in matmul_nbits_common. This follows existing ORT patterns well.
2. Good reuse of existing kernel infrastructure: The RunSkipLayerNormProgram and RunLayerNormProgram refactors extract dispatch logic from the existing kernels so the fused ops can reuse them without duplication. This is the right approach.
3. Thorough pattern matching: The MLP fusion handles both decomposed SiLU (Sigmoid + Mul) and fused QuickGelu(alpha=1.0), covering pre- and post-QuickGeluFusion transformer states.
4. EP-gated registration: Both transformers are restricted to kWebGpuExecutionProvider compatible EPs, preventing any side effects on other backends.
5. Good extensibility comments: The code has clear guidance on how to add new activations (GELU for Gemma-style MLPs).
6. Comprehensive testing: Both fusions include graph structure tests (pattern matching positive/negative cases) and runtime correctness tests (MatchesUnfused*) that compare fused vs unfused outputs on WebGPU EP with 1e-3 tolerance. QKV has 3 runtime tests, MLP has 8 — covering simplified/skip/passthrough/no-bias/QuickGelu variants.

🟡 Suggestions for Improvement

1. WouldApply* helpers should be folded into existing CanApply* functions

The PR introduces 6 new WouldApply* overloads (3 dispatch paths × 2 overloads each) that are thin wrappers adding dispatch preconditions (M threshold, has_weight_idx_indirect, vendor checks) before calling the existing CanApply*. For example:

bool WouldApplyDP4AMatMulNBitsInCurrentDispatch(...) {
  return ((M >= kMinMForTileOptimization && !has_weight_idx_indirect) ||
          y->DataType() == GetType<float>() ||
          vendor == "qualcomm") &&
         CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a);
}

Problems:

• The CanApply* vs WouldApply* distinction is artificial — callers just need "should I use this path?"
• Doubles the API surface (~6 functions, ~100+ lines) for what could be achieved by extending CanApply* with optional parameters
• The verbose naming (WouldApplySubgroupMatrixMatMulNBitsInCurrentDispatch — 52 chars) hurts readability
• The fused ops need WouldApply* to replicate the unfused dispatch decision, so these must stay in sync by definition — arguing for a single function

Suggestion: Extend CanApply* with the dispatch-relevant parameters using defaults:

// Existing callers unchanged (defaults kick in):
bool CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a,
                                    uint32_t M = 0,
                                    bool has_weight_idx_indirect = false);

This eliminates ~6 functions and ~100 lines, keeps a single source of truth for dispatch decisions, and avoids the confusing two-level naming.

2. Significant code duplication between MLP and QKV fusion transformers

The following helper functions are copy-pasted identically between matmul_nbits_mlp_fusion.cc and matmul_nbits_qkv_fusion.cc:

• HasInput()
• IsSupportedSimplifiedLayerNormalization()
• IsSupportedSkipSimplifiedLayerNormalization()
• HasProducedOutput()
• GetIntAttr()
• GetFloatAttr()

These should be extracted into a shared header (e.g., core/optimizer/matmul_nbits_fusion_utils.h). This would reduce ~60 lines of duplication and prevent future divergence if validation logic changes.

3. QKV fusion Q/K/V identification heuristic is scoped to GQA/MQA

In GetQkvNodes(), the code identifies Q vs K/V based solely on the N attribute — it assumes Q has a different N than K and V (which share the same N). This works for GQA/MQA models (Qwen3, Llama3, Mistral) where N_q ≠ N_kv, but won't match classic MHA models (GPT-2, BERT) where all three projections have identical N. The fusion returns nullopt — safe but a missed optimization.

Suggestion: Add a comment clarifying this is intentionally scoped to GQA/MQA, or consider a fallback for equal-N cases.

4. Hardcoded bits == 4 and block_size == 32 constraints in fusion

Both fusion transformers hardcode:

gate_bits == 4 && gate_block_size == 32

But the MLP kernel header claims to support 2-bit and 8-bit:

ORT_ENFORCE(bits_ == 4 || bits_ == 8 || bits_ == 2, ...)

The asymmetry is safe (fusion won't fire for unsupported configs) but should be documented — either narrow the kernel's ORT_ENFORCE to match, or add a comment explaining why the fusion is more restrictive.

📝 Minor Nits

• The norm != nullptr check in preserve_skip_output (MLP fusion) is redundant — norm is guaranteed non-null at that point due to the earlier guard.
• IsGraphOutput in QKV fusion does a linear scan over graph outputs. The MLP fusion already uses graph.NodeProducesGraphOutput() (hash-set backed) — consider using the same approach in QKV for consistency.
• matched_shape uses const char* — could be std::string_view for consistency.

Verdict

This is a well-executed performance optimization with meaningful impact (25-30% decode speedup) and thorough testing (both graph structure and runtime correctness). The primary improvement opportunity is consolidating WouldApply* into the existing CanApply* functions — this would eliminate ~6 functions and ~100 lines while keeping a single source of truth for dispatch decisions. Secondary improvements are extracting duplicated helpers between the two fusion transformers and documenting the GQA/MQA scope limitation of the QKV heuristic. The overall design is architecturally sound and follows established ORT patterns.

Addressed the "Suggestions for improvement" and "Minor nits"

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

Copilot reviewed 31 out of 31 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 31 out of 31 changed files in this pull request and generated 7 comments.