onnxruntime
88ca23fc - [WebGPU] QKV and MLP layer fusions for Qwen3-style models (#28280)

Commit
94 days ago
[WebGPU] QKV and MLP layer fusions for Qwen3-style models (#28280) ### 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 <img width="344" height="140" alt="image" src="https://github.com/user-attachments/assets/0f5d7cfb-05f9-4f25-acb5-4becb8f5addd" /> AFTER (**120+ decode TPS**): PR branch <img width="359" height="134" alt="image" src="https://github.com/user-attachments/assets/f1254d8e-a400-4dbb-9d06-ab6116f929bb" /> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Author
Parents
Loading