sync 91d990483a17

Browse files

Files changed (14) hide show

README.md +26 -16
build/webgpu/bench.json +9 -10
build/webgpu/manifest.json +199 -543
build/webgpu/metadata.json +24 -15
build/webgpu/qmoe-fc1-activation-grouped-sgmat.wgsl.jinja +26 -30
build/webgpu/qmoe-fc1-activation-grouped.wgsl.jinja +1 -1
build/webgpu/qmoe-fc1-activation-stage.wgsl.jinja +4 -5
build/webgpu/qmoe-mix-stage.wgsl.jinja +4 -5
build/webgpu/qmoe-output-grouped-sgmat.wgsl.jinja +28 -34
build/webgpu/qmoe-output-grouped.wgsl.jinja +1 -1
build/webgpu/qmoe-output-stage.wgsl.jinja +4 -5
build/webgpu/qmoe-output-zero-inter.wgsl.jinja +3 -4
build/webgpu/qmoe-route-stage.wgsl.jinja +4 -5
build/webgpu/test.json +199 -22

README.md CHANGED Viewed

@@ -12,26 +12,26 @@ tags:
 ## Description
-Quantized mixture of experts over float32 activations and raw uint8-packed integer expert weights. The package implements 4- and 8-bit symmetric dequantization, ReLU, and interleaved SwiGLU for rank-2 or rank-3 input. Routing uses the softmax over the selected top-k logits, matching ONNX Runtime QMoE when the omitted `router_weights` input is absent. Biases, explicit zero points, FC3, sparse mixing, provider-prepacked layouts, and FP4/FP8 modes are not implemented.
 See the [ONNX Runtime `QMoE` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.QMoE) for the reference semantics.
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `input` | `inputT` | `T` | — | — | Token activations of shape `(num_tokens, hidden_size)` or `(batch_size, sequence_length, hidden_size)`. | required |
-| `router_probs` | `routerT` | `T` | `2` | — | Routing logits of shape `(num_tokens, num_experts)`, where `num_tokens` is the product of the leading input dimensions. | required |
-| `fc1_experts_weights` | `fc1T` | `T1` | `3` | — | Raw packed FC1 weights of shape `(num_experts, fusion_size * inter_size, hidden_size / pack_size)`, where `fusion_size` is 2 only for interleaved SwiGLU and `pack_size` is 2 for 4-bit or 1 for 8-bit weights. | required |
-| `fc1_scales` | `fc1ScalesT` | `T2` | — | — | Required FC1 dequantization scales: rank 2 `(num_experts, fusion_size * inter_size)` for column-wise quantization, or rank 3 with a trailing `hidden_size / block_size` axis. | required |
-| `fc2_experts_weights` | `fc2T` | `T1` | `3` | — | Raw packed FC2 weights of shape `(num_experts, hidden_size, inter_size / pack_size)`. | required |
-| `fc2_scales` | `fc2ScalesT` | `T2` | — | — | Required FC2 dequantization scales: rank 2 `(num_experts, hidden_size)` for column-wise quantization, or rank 3 with a trailing `inter_size / block_size` axis. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `output` | `outputT` | `T` | same as `input` | same as `input` | Routed expert output with the same shape as `input`. | required |
 ## Attributes
@@ -42,15 +42,15 @@ Attributes and default values (overridable per request):
 | `activation_alpha` | `1` | Alpha used by SwiGLU; the exact standard default is 1. |
 | `activation_beta` | `0` | Beta added to the SwiGLU linear branch; the exact standard default is 0. |
 | `activation_type` | `"relu"` | Activation applied after FC1. This package supports `relu` and `swiglu`; the exact standard default is `relu`. |
 | `expert_weight_bits` | `4` | Integer expert-weight bit width. This package supports 4 and 8; the exact standard default is 4. |
 | `k` | `1` | Number of experts selected per token; the exact standard default is 1. |
-| `normalize_routing_weights` | `0` | Accepted values are 0 and 1. With the separate `router_weights` input omitted, ONNX Runtime QMoE applies a softmax over the selected top-k logits for either value; the attribute only distinguishes the unsupported separate-weight path. |
 | `quant_type` | `"int"` | Quantization family. This package supports only the exact standard default `int`. |
 | `swiglu_fusion` | `0` | SwiGLU packing mode. ReLU uses the exact standard default 0; supported SwiGLU interleaves gate/up FC1 rows with value 1. |
 | `use_sparse_mixer` | `0` | Whether to use sparse-mixer routing. The exact standard default and only supported value is 0. |
 | `weights_prepacked` | `-1` | Provider weight-layout selector. Values -1 and 0 both consume the public raw packed tensor layout; provider-specific prepacked value 1 is not portable and is rejected. |
-| `block_size` | — | Optional quantization block size along the reduction dimension. Omission selects column-wise scaling; a supplied value must be a power of two of at least 16 and divide both `hidden_size` and `inter_size`. |
-| `swiglu_limit` | — | Optional SwiGLU clamp limit. Omission means no finite clamp. |
 ## Type constraints
@@ -60,13 +60,19 @@ Attributes and default values (overridable per request):
 | `T1` | `uint8` |
 | `T2` | `float32` |
 ## Device requirements
 Some implementation variants require `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -85,10 +91,14 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Description
+Quantized mixture of experts over float32 activations and raw uint8-packed integer expert weights. The package implements 4- and 8-bit symmetric dequantization, ReLU, and interleaved SwiGLU for rank-2 or rank-3 input. Routing uses the softmax over the selected top-k logits with `router_weights` absent. Biases, explicit zero points, FC3, sparse mixing, provider-prepacked layouts, and FP4/FP8 modes are not implemented.
 See the [ONNX Runtime `QMoE` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.QMoE) for the reference semantics.
 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `inputT` | `input` | `T` | — | — | Token activations of shape `(num_tokens, hidden_size)` or `(batch_size, sequence_length, hidden_size)`. | required |
+| `routerT` | `router_probs` | `T` | `2` | — | Routing logits of shape `(num_tokens, num_experts)`, where `num_tokens` is the product of the leading input dimensions. | required |
+| `fc1T` | `fc1_experts_weights` | `T1` | `3` | — | Raw packed FC1 weights of shape `(num_experts, fusion_size * inter_size, hidden_size / pack_size)`, where `fusion_size` is 2 only for interleaved SwiGLU and `pack_size` is 2 for 4-bit or 1 for 8-bit weights. | required |
+| `fc1ScalesT` | `fc1_scales` | `T2` | — | — | Required FC1 dequantization scales: rank 2 `(num_experts, fusion_size * inter_size)` for column-wise quantization, or rank 3 with a trailing `hidden_size / block_size` axis. | required |
+| `fc2T` | `fc2_experts_weights` | `T1` | `3` | — | Raw packed FC2 weights of shape `(num_experts, hidden_size, inter_size / pack_size)`. | required |
+| `fc2ScalesT` | `fc2_scales` | `T2` | — | — | Required FC2 dequantization scales: rank 2 `(num_experts, hidden_size)` for column-wise quantization, or rank 3 with a trailing `inter_size / block_size` axis. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `outputT` | `output` | `T` | same as `inputT` | same as `inputT` | Routed expert output with the same shape as `input`. | required |
 ## Attributes
 | `activation_alpha` | `1` | Alpha used by SwiGLU; the exact standard default is 1. |
 | `activation_beta` | `0` | Beta added to the SwiGLU linear branch; the exact standard default is 0. |
 | `activation_type` | `"relu"` | Activation applied after FC1. This package supports `relu` and `swiglu`; the exact standard default is `relu`. |
+| `block_size` | — | Optional quantization block size along the reduction dimension. Omission selects column-wise scaling; a supplied value must be a power of two of at least 16 and divide both `hidden_size` and `inter_size`. |
 | `expert_weight_bits` | `4` | Integer expert-weight bit width. This package supports 4 and 8; the exact standard default is 4. |
 | `k` | `1` | Number of experts selected per token; the exact standard default is 1. |
+| `normalize_routing_weights` | `0` | Accepted values are 0 and 1. With the separate `router_weights` input omitted, routing applies a softmax over the selected top-k logits for either value; the attribute only distinguishes the unsupported separate-weight path. |
 | `quant_type` | `"int"` | Quantization family. This package supports only the exact standard default `int`. |
 | `swiglu_fusion` | `0` | SwiGLU packing mode. ReLU uses the exact standard default 0; supported SwiGLU interleaves gate/up FC1 rows with value 1. |
+| `swiglu_limit` | — | Optional SwiGLU clamp limit. Omission means no finite clamp. |
 | `use_sparse_mixer` | `0` | Whether to use sparse-mixer routing. The exact standard default and only supported value is 0. |
 | `weights_prepacked` | `-1` | Provider weight-layout selector. Values -1 and 0 both consume the public raw packed tensor layout; provider-specific prepacked value 1 is not portable and is rejected. |
 ## Type constraints
 | `T1` | `uint8` |
 | `T2` | `float32` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `quant_grouped_sgmat_routed` — Expert-grouped f32 matrix multiplication with per-invocation routing slots and reuse of the input tile for output publication. The larger operand/output footprint determines workgroup storage admission from device limits.
 ## Device requirements
 Some implementation variants require `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.QMoE",
   "tunableSpace": {
     "workgroupSize": [64, 128, 256],
     "groupThreads": [8, 16],
@@ -109,7 +108,7 @@
         "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
       },
       "provenance": {
-        "notes": "Scored as bandwidth, not gflops: a decode step reads whole expert weight matrices to produce one output row, so memory is its roofline and the compute peak is the wrong denominator. The prefill cases stay on gflops. The bytes are GPU storage, not the packed blob: a uint8 tensor takes one byte per u32 element here, so a b-bit code costs b/2 bytes of traffic."
       }
     },
     {
@@ -144,7 +143,7 @@
         "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
       },
       "provenance": {
-        "notes": "Scored as bandwidth, not gflops: a decode step reads whole expert weight matrices to produce one output row, so memory is its roofline and the compute peak is the wrong denominator. The prefill cases stay on gflops. The bytes are GPU storage, not the packed blob: a uint8 tensor takes one byte per u32 element here, so a b-bit code costs b/2 bytes of traffic."
       }
     },
     {
@@ -180,7 +179,7 @@
         "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
       },
       "provenance": {
-        "notes": "Scored as bandwidth, not gflops: a decode step reads whole expert weight matrices to produce one output row, so memory is its roofline and the compute peak is the wrong denominator. The prefill cases stay on gflops. The bytes are GPU storage, not the packed blob: a uint8 tensor takes one byte per u32 element here, so a b-bit code costs b/2 bytes of traffic."
       }
     },
     {
@@ -215,7 +214,7 @@
         "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
       },
       "provenance": {
-        "notes": "Scored as bandwidth, not gflops: a decode step reads whole expert weight matrices to produce one output row, so memory is its roofline and the compute peak is the wrong denominator. The prefill cases stay on gflops. The bytes are GPU storage, not the packed blob: a uint8 tensor takes one byte per u32 element here, so a b-bit code costs b/2 bytes of traffic."
       }
     },
     {
@@ -249,7 +248,7 @@
         "metrics": [{ "type": "bandwidth", "value": "args.topK * 2 * args.hidden * args.inter * args.bits / 2" }]
       },
       "provenance": {
-        "notes": "Scored as bandwidth, not gflops: a decode step reads whole expert weight matrices to produce one output row, so memory is its roofline and the compute peak is the wrong denominator. The prefill cases stay on gflops. The bytes are GPU storage, not the packed blob: a uint8 tensor takes one byte per u32 element here, so a b-bit code costs b/2 bytes of traffic."
       }
     },
     {
@@ -419,7 +418,7 @@
       "name": "qmoe-q4-qwen3-moe-decode-t1-h2048-e64-k8-i768",
       "preset": "model",
       "provenance": {
-        "notes": "Qwen3-MoE class defaults (hidden_size 2048, moe_intermediate_size 768, num_experts_per_tok 8); expert count capped at 64 of the config's 128 to keep each weight tensor under 512 MB. Scored as bandwidth, not gflops: a decode step reads whole expert weight matrices to produce one output row, so memory is its roofline and the compute peak is the wrong denominator. The prefill cases stay on gflops. The bytes are GPU storage, not the packed blob: a uint8 tensor takes one byte per u32 element here, so a b-bit code costs b/2 bytes of traffic."
       },
       "vars": { "tokens": 1, "hidden": 2048, "experts": 64, "inter": 768, "topK": 8, "bits": 4 },
       "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 8 },
@@ -489,7 +488,7 @@
       "name": "qmoe-q4-olmoe-decode-t1-h2048-e32-k8-i2048",
       "preset": "model",
       "provenance": {
-        "notes": "OLMoE class defaults (hidden_size 2048, intermediate_size 2048, num_experts_per_tok 8); expert count capped at 32 of the config's 64. Scored as bandwidth, not gflops: a decode step reads whole expert weight matrices to produce one output row, so memory is its roofline and the compute peak is the wrong denominator. The prefill cases stay on gflops. The bytes are GPU storage, not the packed blob: a uint8 tensor takes one byte per u32 element here, so a b-bit code costs b/2 bytes of traffic."
       },
       "vars": { "tokens": 1, "hidden": 2048, "experts": 32, "inter": 2048, "topK": 8, "bits": 4 },
       "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 8 },
@@ -557,7 +556,7 @@
       "name": "qmoe-q4-gpt-oss-decode-t1-h2880-e16-k4-i2880",
       "preset": "model",
       "provenance": {
-        "notes": "GPT-OSS class defaults (hidden_size 2880, intermediate_size 2880, num_experts_per_tok 4); expert count capped at 16 of the config's 128. Scored as bandwidth, not gflops: a decode step reads whole expert weight matrices to produce one output row, so memory is its roofline and the compute peak is the wrong denominator. The prefill cases stay on gflops. The bytes are GPU storage, not the packed blob: a uint8 tensor takes one byte per u32 element here, so a b-bit code costs b/2 bytes of traffic."
       },
       "vars": { "tokens": 1, "hidden": 2880, "experts": 16, "inter": 2880, "topK": 4, "bits": 4 },
       "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 4 },
@@ -625,7 +624,7 @@
       "name": "qmoe-q4-phimoe-decode-t1-h4096-e4-k2-i6400",
       "preset": "model",
       "provenance": {
-        "notes": "Phi-MoE class defaults (hidden_size 4096, intermediate_size 6400, num_experts_per_tok 2); expert count capped at 4 of the config's 16 -- the widest expert in this set, so it is the case that stresses the reduction rather than the routing. Scored as bandwidth, not gflops: a decode step reads whole expert weight matrices to produce one output row, so memory is its roofline and the compute peak is the wrong denominator. The prefill cases stay on gflops. The bytes are GPU storage, not the packed blob: a uint8 tensor takes one byte per u32 element here, so a b-bit code costs b/2 bytes of traffic."
       },
       "vars": { "tokens": 1, "hidden": 4096, "experts": 4, "inter": 6400, "topK": 2, "bits": 4 },
       "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 2 },

 {
   "tunableSpace": {
     "workgroupSize": [64, 128, 256],
     "groupThreads": [8, 16],
         "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
       },
       "provenance": {
+        "notes": "The bandwidth metric counts the selected experts' three quantized weight matrices at `bits/2` bytes per logical weight in the bound storage layout."
       }
     },
     {
         "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
       },
       "provenance": {
+        "notes": "The bandwidth metric counts the selected experts' three quantized weight matrices at `bits/2` bytes per logical weight in the bound storage layout."
       }
     },
     {
         "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
       },
       "provenance": {
+        "notes": "The bandwidth metric counts the selected experts' three quantized weight matrices at `bits/2` bytes per logical weight in the bound storage layout."
       }
     },
     {
         "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
       },
       "provenance": {
+        "notes": "The bandwidth metric counts the selected experts' three quantized weight matrices at `bits/2` bytes per logical weight in the bound storage layout."
       }
     },
     {
         "metrics": [{ "type": "bandwidth", "value": "args.topK * 2 * args.hidden * args.inter * args.bits / 2" }]
       },
       "provenance": {
+        "notes": "The bandwidth metric counts the selected expert's two quantized weight matrices at `bits/2` bytes per logical weight in the bound storage layout."
       }
     },
     {
       "name": "qmoe-q4-qwen3-moe-decode-t1-h2048-e64-k8-i768",
       "preset": "model",
       "provenance": {
+        "notes": "Qwen3-MoE geometry uses hidden size 2,048, intermediate size 768, top-k 8, and 64 instantiated experts to keep each weight tensor below 512 MB. The bandwidth metric counts three selected quantized matrices at `bits/2` bytes per logical weight."
       },
       "vars": { "tokens": 1, "hidden": 2048, "experts": 64, "inter": 768, "topK": 8, "bits": 4 },
       "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 8 },
       "name": "qmoe-q4-olmoe-decode-t1-h2048-e32-k8-i2048",
       "preset": "model",
       "provenance": {
+        "notes": "OLMoE geometry uses hidden and intermediate sizes of 2,048, top-k 8, and 32 instantiated experts. The bandwidth metric counts three selected quantized matrices at `bits/2` bytes per logical weight."
       },
       "vars": { "tokens": 1, "hidden": 2048, "experts": 32, "inter": 2048, "topK": 8, "bits": 4 },
       "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 8 },
       "name": "qmoe-q4-gpt-oss-decode-t1-h2880-e16-k4-i2880",
       "preset": "model",
       "provenance": {
+        "notes": "GPT-OSS geometry uses hidden and intermediate sizes of 2,880, top-k 4, and 16 instantiated experts. The bandwidth metric counts three selected quantized matrices at `bits/2` bytes per logical weight."
       },
       "vars": { "tokens": 1, "hidden": 2880, "experts": 16, "inter": 2880, "topK": 4, "bits": 4 },
       "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 4 },
       "name": "qmoe-q4-phimoe-decode-t1-h4096-e4-k2-i6400",
       "preset": "model",
       "provenance": {
+        "notes": "Phi-MoE geometry uses hidden size 4,096, intermediate size 6,400, top-k 2, and four instantiated experts. The bandwidth metric counts three selected quantized matrices at `bits/2` bytes per logical weight."
       },
       "vars": { "tokens": 1, "hidden": 4096, "experts": 4, "inter": 6400, "topK": 2, "bits": 4 },
       "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 2 },

build/webgpu/manifest.json CHANGED Viewed

@@ -2,75 +2,28 @@
   "domain": "com.microsoft",
   "name": "QMoE",
   "sinceVersion": 1,
-  "description": "Quantized mixture of experts over float32 activations and raw uint8-packed integer expert weights. The package implements 4- and 8-bit symmetric dequantization, ReLU, and interleaved SwiGLU for rank-2 or rank-3 input. Routing uses the softmax over the selected top-k logits, matching ONNX Runtime QMoE when the omitted `router_weights` input is absent. Biases, explicit zero points, FC3, sparse mixing, provider-prepacked layouts, and FP4/FP8 modes are not implemented.",
-  "inputs": [
-    {
-      "role": "input",
-      "dtype": "T",
-      "description": "Token activations of shape `(num_tokens, hidden_size)` or `(batch_size, sequence_length, hidden_size)`."
-    },
-    {
-      "role": "router_probs",
-      "dtype": "T",
-      "rank": 2,
-      "description": "Routing logits of shape `(num_tokens, num_experts)`, where `num_tokens` is the product of the leading input dimensions."
-    },
-    {
-      "role": "fc1_experts_weights",
-      "dtype": "T1",
-      "rank": 3,
-      "description": "Raw packed FC1 weights of shape `(num_experts, fusion_size * inter_size, hidden_size / pack_size)`, where `fusion_size` is 2 only for interleaved SwiGLU and `pack_size` is 2 for 4-bit or 1 for 8-bit weights."
-    },
-    {
-      "role": "fc1_scales",
-      "dtype": "T2",
-      "description": "Required FC1 dequantization scales: rank 2 `(num_experts, fusion_size * inter_size)` for column-wise quantization, or rank 3 with a trailing `hidden_size / block_size` axis."
-    },
-    {
-      "role": "fc2_experts_weights",
-      "dtype": "T1",
-      "rank": 3,
-      "description": "Raw packed FC2 weights of shape `(num_experts, hidden_size, inter_size / pack_size)`."
-    },
-    {
-      "role": "fc2_scales",
-      "dtype": "T2",
-      "description": "Required FC2 dequantization scales: rank 2 `(num_experts, hidden_size)` for column-wise quantization, or rank 3 with a trailing `inter_size / block_size` axis."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "output",
-      "dtype": "T",
-      "shape": "shapes.input",
-      "description": "Routed expert output with the same shape as `input`."
-    }
-  ],
-  "attributes": {
-    "activation_alpha": 1,
-    "activation_beta": 0,
-    "activation_type": "relu",
-    "expert_weight_bits": 4,
-    "k": 1,
-    "normalize_routing_weights": 0,
-    "quant_type": "int",
-    "swiglu_fusion": 0,
-    "use_sparse_mixer": 0,
-    "weights_prepacked": -1
   },
-  "attributeDescriptions": {
-    "activation_alpha": "Alpha used by SwiGLU; the exact standard default is 1.",
-    "activation_beta": "Beta added to the SwiGLU linear branch; the exact standard default is 0.",
-    "activation_type": "Activation applied after FC1. This package supports `relu` and `swiglu`; the exact standard default is `relu`.",
-    "block_size": "Optional quantization block size along the reduction dimension. Omission selects column-wise scaling; a supplied value must be a power of two of at least 16 and divide both `hidden_size` and `inter_size`.",
-    "expert_weight_bits": "Integer expert-weight bit width. This package supports 4 and 8; the exact standard default is 4.",
-    "k": "Number of experts selected per token; the exact standard default is 1.",
-    "normalize_routing_weights": "Accepted values are 0 and 1. With the separate `router_weights` input omitted, ONNX Runtime QMoE applies a softmax over the selected top-k logits for either value; the attribute only distinguishes the unsupported separate-weight path.",
-    "quant_type": "Quantization family. This package supports only the exact standard default `int`.",
-    "swiglu_fusion": "SwiGLU packing mode. ReLU uses the exact standard default 0; supported SwiGLU interleaves gate/up FC1 rows with value 1.",
-    "swiglu_limit": "Optional SwiGLU clamp limit. Omission means no finite clamp.",
-    "use_sparse_mixer": "Whether to use sparse-mixer routing. The exact standard default and only supported value is 0.",
-    "weights_prepacked": "Provider weight-layout selector. Values -1 and 0 both consume the public raw packed tensor layout; provider-specific prepacked value 1 is not portable and is rejected."
   },
   "attributeConstraints": {
     "activation_type": { "values": ["relu", "swiglu"] },
@@ -82,25 +35,16 @@
     "weights_prepacked": { "values": [-1, 0] }
   },
   "typeConstraints": { "T": ["float32"], "T1": ["uint8"], "T2": ["float32"] },
-  "args": {
-    "inputT": { "kind": "tensor", "semantic": "input", "role": "input" },
-    "routerT": { "kind": "tensor", "semantic": "router_probs", "role": "input" },
-    "fc1T": { "kind": "tensor", "semantic": "fc1_experts_weights", "role": "input" },
-    "fc1ScalesT": { "kind": "tensor", "semantic": "fc1_scales", "role": "input" },
-    "fc2T": { "kind": "tensor", "semantic": "fc2_experts_weights", "role": "input" },
-    "fc2ScalesT": { "kind": "tensor", "semantic": "fc2_scales", "role": "input" },
-    "outputT": { "kind": "tensor", "semantic": "output", "role": "output" }
-  },
   "tunables": {
-    "workgroupSize": 64,
-    "decodeLanes": 32,
-    "decodeBlockTarget": 1024,
-    "decodeMinLaneTrips": 4,
-    "groupThreads": 8,
-    "groupRegM": 4,
-    "groupRegN": 4,
-    "groupTileK": 16,
-    "groupRouteWorkgroup": 256
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
@@ -108,28 +52,28 @@
     "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
     "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
     "wave32Effective": "wave32Adapter or pinSubgroupSize32",
-    "hiddenSize": "dim(shapes.input, ranks.input - 1)",
-    "numTokens": "numel(shapes.input) / max(1, hiddenSize)",
     "topK": "attrs.k",
     "weightBits": "attrs.expert_weight_bits",
     "quantBlockSize": "attrs.block_size if has(attrs, \"block_size\") else 0",
     "packSize": "2 if weightBits == 4 else 1",
     "quantMidpoint": "8 if weightBits == 4 else 128",
     "fusionSize": "2 if attrs.activation_type == \"swiglu\" and attrs.swiglu_fusion == 1 else 1",
-    "interSize": "dim(shapes.fc1_experts_weights, 1) / fusionSize",
-    "fc1PackedCols": "dim(shapes.fc1_experts_weights, 2)",
-    "fc2PackedCols": "dim(shapes.fc2_experts_weights, 2)",
     "colWiseScales": "quantBlockSize == 0",
     "fc1ScaleBlocks": "1 if colWiseScales else hiddenSize / max(1, quantBlockSize)",
     "fc2ScaleBlocks": "1 if colWiseScales else interSize / max(1, quantBlockSize)",
     "activationSupported": "(attrs.activation_type == \"relu\" and attrs.swiglu_fusion == 0) or (attrs.activation_type == \"swiglu\" and attrs.swiglu_fusion == 1)",
     "routingModeSupported": "attrs.normalize_routing_weights == 0 or attrs.normalize_routing_weights == 1",
     "rawWeightLayout": "attrs.weights_prepacked == -1 or attrs.weights_prepacked == 0",
-    "inputOutputShapeOk": "((ranks.input == 2 and ranks.output == 2 and dim(shapes.output, 0) == dim(shapes.input, 0) and dim(shapes.output, 1) == dim(shapes.input, 1)) or (ranks.input == 3 and ranks.output == 3 and dim(shapes.output, 0) == dim(shapes.input, 0) and dim(shapes.output, 1) == dim(shapes.input, 1) and dim(shapes.output, 2) == dim(shapes.input, 2))) and hiddenSize > 0",
     "quantBlockSizeOk": "colWiseScales or (quantBlockSize >= 16 and pow2ceil(quantBlockSize) == quantBlockSize and hiddenSize % quantBlockSize == 0 and interSize % quantBlockSize == 0)",
-    "quantScalesOk": "tensorDtypes.fc1_scales == \"float32\" and tensorDtypes.fc2_scales == \"float32\" and dim(shapes.fc1_scales, 0) == dim(shapes.router_probs, 1) and dim(shapes.fc2_scales, 0) == dim(shapes.router_probs, 1) and dim(shapes.fc1_scales, 1) == dim(shapes.fc1_experts_weights, 1) and dim(shapes.fc2_scales, 1) == hiddenSize and ((ranks.fc1_scales == 2 and ranks.fc2_scales == 2) if colWiseScales else (ranks.fc1_scales == 3 and ranks.fc2_scales == 3 and dim(shapes.fc1_scales, 2) == fc1ScaleBlocks and dim(shapes.fc2_scales, 2) == fc2ScaleBlocks))",
-    "quantShapeOk": "inputOutputShapeOk and ranks.router_probs == 2 and ranks.fc1_experts_weights == 3 and ranks.fc2_experts_weights == 3 and dim(shapes.router_probs, 0) == numTokens and dim(shapes.fc1_experts_weights, 0) == dim(shapes.router_probs, 1) and dim(shapes.fc2_experts_weights, 0) == dim(shapes.router_probs, 1) and dim(shapes.fc1_experts_weights, 1) % fusionSize == 0 and dim(shapes.fc2_experts_weights, 1) == hiddenSize and dim(shapes.fc1_experts_weights, 2) * packSize == hiddenSize and dim(shapes.fc2_experts_weights, 2) * packSize == interSize and quantBlockSizeOk and quantScalesOk",
-    "quantContract": "activationSupported and routingModeSupported and rawWeightLayout and quantShapeOk and topK >= 1 and topK <= dim(shapes.router_probs, 1)",
     "workgroupSizeOk": "tunables.workgroupSize >= 1 and tunables.workgroupSize <= deviceWorkgroupCap",
     "hiddenChunkFits": "topK * interSize * 4 <= device.limits.maxStorageBufferBindingSize and topK * interSize * 4 <= device.limits.maxBufferSize",
     "hiddenChunkTokens": "numTokens if interSize == 0 else min(numTokens, max(1, floor(min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) / (topK * interSize * 4))))",
@@ -141,421 +85,89 @@
     "groupTileKVec": "ceilDiv(tunables.groupTileK, 4)",
     "groupThreadCount": "tunables.groupThreads * tunables.groupThreads",
     "groupSlots": "hiddenChunkTokens * topK",
-    "groupMaxTiles": "ceilDiv(groupSlots, max(1, groupTileM)) + dim(shapes.router_probs, 1)",
     "groupSharedBytes": "(groupTileM * tunables.groupTileK + 2 * groupTileN * tunables.groupTileK + groupTileM) * 4",
     "groupSlotOutBytes": "hiddenChunkTokens * topK * hiddenSize * 4",
-    "groupedDeviceOk": "groupThreadCount <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.groupThreads <= device.limits.maxComputeWorkgroupSizeX and tunables.groupThreads <= device.limits.maxComputeWorkgroupSizeY and tunables.groupRouteWorkgroup <= deviceWorkgroupCap and groupSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.router_probs, 1) * 8 <= device.limits.maxComputeWorkgroupStorageSize and tunables.groupTileK % 4 == 0",
-    "groupedDispatchOk": "groupMaxTiles <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(interSize, max(1, groupTileN)) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(hiddenSize, max(1, groupTileN)) <= device.limits.maxComputeWorkgroupsPerDimension",
     "groupSlotOutFits": "groupSlotOutBytes <= device.limits.maxStorageBufferBindingSize and groupSlotOutBytes <= device.limits.maxBufferSize",
-    "groupedShapeOk": "numTokens * topK * 4 >= groupTileM * dim(shapes.router_probs, 1)",
     "groupedContract": "quantContract and workgroupSizeOk and hiddenChunkFits and routedScratchFits and interSize > 0 and groupedDeviceOk and groupedDispatchOk and groupSlotOutFits and groupedShapeOk",
     "sgmatWorkgroup": "128",
     "sgmatSubgroups": "4",
     "sgmatRowSubtiles": "2",
-    "groupedSgmatSharedBytes": "groupTileM * 4 + groupTileM * 32 * 4 + 64 * 32 * 4 + sgmatSubgroups * 4 * 64 * 4",
-    "groupedSgmatOk": "groupedContract and groupTileM == 32 and interSize % 32 == 0 and groupedSgmatSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(hiddenSize, 64) <= device.limits.maxComputeWorkgroupsPerDimension and hiddenSize % 32 == 0",
     "decodeLanesOk": "tunables.decodeLanes >= 1 and tunables.decodeLanes <= tunables.workgroupSize and tunables.workgroupSize % tunables.decodeLanes == 0",
     "decodeRows": "max(1, tunables.workgroupSize / max(1, tunables.decodeLanes))",
     "decodeDeviceOk": "workgroupSizeOk and decodeLanesOk and tunables.decodeLanes <= device.limits.maxComputeWorkgroupSizeX and decodeRows <= device.limits.maxComputeWorkgroupSizeY and tunables.workgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
-    "decodeDispatchOk": "ceilDiv(interSize, decodeRows) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(hiddenSize, decodeRows) <= device.limits.maxComputeWorkgroupsPerDimension and hiddenChunkTokens * topK <= device.limits.maxComputeWorkgroupsPerDimension",
     "decodeLaneDepth": "min(hiddenSize, interSize) / max(1, tunables.decodeLanes)",
     "decodeDepthOk": "decodeLaneDepth >= tunables.decodeMinLaneTrips",
     "splitFc1Blocks": "ceilDiv(numTokens * topK * interSize, max(1, tunables.workgroupSize))",
     "decodeOccupancyOk": "splitFc1Blocks <= tunables.decodeBlockTarget",
-    "decodeContract": "quantContract and decodeDeviceOk and decodeDispatchOk and hiddenChunkFits and routedScratchFits and interSize > 0 and decodeDepthOk and (decodeOccupancyOk or not groupedContract)"
-  },
-  "constants": {
     "tokens": "numTokens",
     "hidden": "hiddenSize",
-    "experts": "dim(shapes.router_probs, 1)",
-    "fc1Rows": "dim(shapes.fc1_experts_weights, 1)",
     "inter": "interSize",
-    "topK": "topK",
     "workgroupSize": "tunables.workgroupSize",
     "activationType": "attrs.activation_type",
     "hasSwigluLimit": "has(attrs, \"swiglu_limit\")",
-    "swigluLimit": "attrs.swiglu_limit if has(attrs, \"swiglu_limit\") else 0",
-    "quantMidpoint": "quantMidpoint",
-    "quantBlockSize": "quantBlockSize",
-    "packSize": "packSize",
-    "fc1PackedCols": "fc1PackedCols",
-    "fc2PackedCols": "fc2PackedCols",
-    "colWiseScales": "colWiseScales",
-    "fc1ScaleBlocks": "fc1ScaleBlocks",
-    "fc2ScaleBlocks": "fc2ScaleBlocks"
   },
-  "bindingSets": {
-    "route": [
-      {
-        "name": "router_probs",
-        "arg": "routerT",
-        "semantic": "router_probs",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "route_expert", "semantic": "routeExpert", "buffer": { "type": "storage" }, "elementType": "u32" },
-      { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "storage" }, "elementType": "f32" }
-    ],
-    "fc1Activation": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "route_expert",
-        "semantic": "routeExpert",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc1_experts_weights",
-        "arg": "fc1T",
-        "semantic": "fc1_experts_weights",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc1_scales",
-        "arg": "fc1ScalesT",
-        "semantic": "fc1_scales",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" },
-            { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" },
-            { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" },
-            {
-              "name": "tokenCount",
-              "type": "u32",
-              "value": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens)"
-            }
-          ]
-        }
-      }
-    ],
-    "output": [
-      { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "route_expert",
-        "semantic": "routeExpert",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "fc2_experts_weights",
-        "arg": "fc2T",
-        "semantic": "fc2_experts_weights",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc2_scales",
-        "arg": "fc2ScalesT",
-        "semantic": "fc2_scales",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" },
-            {
-              "name": "tokenCount",
-              "type": "u32",
-              "value": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens)"
-            }
-          ]
-        }
-      }
-    ],
-    "zeroInterOutput": [
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "f32"
-      }
-    ],
-    "group": [
-      {
-        "name": "route_expert",
-        "semantic": "routeExpert",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "storage" }, "elementType": "u32" },
-      { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" },
-            {
-              "name": "tokenCount",
-              "type": "u32",
-              "value": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens)"
-            }
-          ]
-        }
-      }
-    ],
-    "fc1Grouped": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
-      { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
-      {
-        "name": "fc1_experts_weights",
-        "arg": "fc1T",
-        "semantic": "fc1_experts_weights",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc1_scales",
-        "arg": "fc1ScalesT",
-        "semantic": "fc1_scales",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" },
-            { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" },
-            { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }
-          ]
-        }
-      }
-    ],
-    "outputGrouped": [
-      { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
-      { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
-      {
-        "name": "fc2_experts_weights",
-        "arg": "fc2T",
-        "semantic": "fc2_experts_weights",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc2_scales",
-        "arg": "fc2ScalesT",
-        "semantic": "fc2_scales",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "slot_out", "semantic": "slotOut", "buffer": { "type": "storage" }, "elementType": "f32" }
-    ],
-    "mix": [
-      { "name": "slot_out", "semantic": "slotOut", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" },
-            {
-              "name": "tokenCount",
-              "type": "u32",
-              "value": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens)"
-            }
-          ]
-        }
-      }
-    ],
-    "fc1ActivationIo": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "route_expert",
-        "semantic": "routeExpert",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc1_experts_weights",
-        "arg": "fc1T",
-        "semantic": "fc1_experts_weights",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc1_scales",
-        "arg": "fc1ScalesT",
-        "semantic": "fc1_scales",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }
-    ],
-    "fc1ActivationScheduled": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "route_expert",
-        "semantic": "routeExpert",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc1_experts_weights",
-        "arg": "fc1T",
-        "semantic": "fc1_experts_weights",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc1_scales",
-        "arg": "fc1ScalesT",
-        "semantic": "fc1_scales",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" },
-            { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" },
-            { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }
-          ]
-        }
-      }
-    ],
-    "outputIo": [
-      { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "route_expert",
-        "semantic": "routeExpert",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "fc2_experts_weights",
-        "arg": "fc2T",
-        "semantic": "fc2_experts_weights",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc2_scales",
-        "arg": "fc2ScalesT",
-        "semantic": "fc2_scales",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "f32"
-      }
-    ],
-    "outputScheduled": [
-      { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "route_expert",
-        "semantic": "routeExpert",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "fc2_experts_weights",
-        "arg": "fc2T",
-        "semantic": "fc2_experts_weights",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "fc2_scales",
-        "arg": "fc2ScalesT",
-        "semantic": "fc2_scales",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [{ "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }]
         }
-      }
-    ]
   },
   "variants": [
     {
@@ -566,12 +178,14 @@
         {
           "id": "output_stage",
           "name": "QMoE.OutputStageZeroInter",
-          "source": {
-            "shader": "qmoe-output-zero-inter.wgsl.jinja",
-            "inputs": { "outputElementCount": "numel(shapes.output)" }
-          },
-          "bindings": "zeroInterOutput",
-          "dispatch": { "threads": "numel(shapes.output)", "workgroupSize": "constants.workgroupSize" }
         }
       ]
     },
@@ -579,13 +193,11 @@
       "id": "quant_grouped_sgmat_routed",
       "priority": 32,
       "when": ["groupedSgmatOk", "wave32Effective"],
-      "constants": {
-        "groupTileM": "groupTileM",
-        "groupRouteWorkgroup": "tunables.groupRouteWorkgroup",
-        "sgmatWorkgroup": "sgmatWorkgroup",
-        "sgmatSubgroups": "sgmatSubgroups",
-        "sgmatRowSubtiles": "sgmatRowSubtiles"
       },
       "intermediates": [
         { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" },
         { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" },
@@ -599,8 +211,12 @@
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
-          "bindings": "route",
-          "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" }
         },
         {
           "id": "token_chunks",
@@ -610,54 +226,48 @@
               "id": "group_stage",
               "name": "QMoE.GroupStage",
               "shader": "expert-group-slots.wgsl.jinja",
-              "bindings": "group",
               "dispatch": {
-                "threads": "constants.groupRouteWorkgroup",
-                "workgroupSize": "constants.groupRouteWorkgroup"
               }
             },
             {
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGroupedSgmat",
               "shader": "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja",
-              "bindings": "fc1Grouped",
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize * fusionSize, 64)" }
             },
             {
               "id": "output_stage",
               "name": "QMoE.OutputStageGroupedSgmat",
               "shader": "qmoe-output-grouped-sgmat.wgsl.jinja",
-              "bindings": "outputGrouped",
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenSize, 64)" }
             },
             {
               "id": "mix_stage",
               "name": "QMoE.MixStage",
               "shader": "qmoe-mix-stage.wgsl.jinja",
-              "bindings": "mix",
               "dispatch": {
-                "threads": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize",
-                "workgroupSize": "constants.workgroupSize"
               }
             }
           ]
         }
-      ],
-      "requires": {
-        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
-        "limits": { "maxComputeWorkgroupStorageSize": 16512 },
-        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
-      }
     },
     {
       "id": "quant_grouped_routed",
       "priority": 30,
       "when": ["groupedContract"],
-      "constants": {
-        "groupTileM": "groupTileM",
-        "groupTileN": "groupTileN",
         "groupTileK": "tunables.groupTileK",
-        "groupTileKVec": "groupTileKVec",
         "groupThreads": "tunables.groupThreads",
         "regM": "tunables.groupRegM",
         "regN": "tunables.groupRegN",
@@ -676,8 +286,12 @@
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
-          "bindings": "route",
-          "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" }
         },
         {
           "id": "token_chunks",
@@ -687,34 +301,36 @@
               "id": "group_stage",
               "name": "QMoE.GroupStage",
               "shader": "expert-group-slots.wgsl.jinja",
-              "bindings": "group",
               "dispatch": {
-                "threads": "constants.groupRouteWorkgroup",
-                "workgroupSize": "constants.groupRouteWorkgroup"
               }
             },
             {
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGrouped",
               "shader": "qmoe-fc1-activation-grouped.wgsl.jinja",
-              "bindings": "fc1Grouped",
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" }
             },
             {
               "id": "output_stage",
               "name": "QMoE.OutputStageGrouped",
               "shader": "qmoe-output-grouped.wgsl.jinja",
-              "bindings": "outputGrouped",
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenSize, groupTileN)" }
             },
             {
               "id": "mix_stage",
               "name": "QMoE.MixStage",
               "shader": "qmoe-mix-stage.wgsl.jinja",
-              "bindings": "mix",
               "dispatch": {
-                "threads": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize",
-                "workgroupSize": "constants.workgroupSize"
               }
             }
           ]
@@ -725,7 +341,7 @@
       "id": "quant_gemv_routed",
       "priority": 20,
       "when": ["decodeContract"],
-      "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" },
       "intermediates": [
         { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" },
         { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" },
@@ -736,8 +352,12 @@
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
-          "bindings": "route",
-          "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" }
         },
         {
           "id": "token_chunks",
@@ -747,7 +367,7 @@
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGemv",
               "shader": "qmoe-fc1-activation-gemv.wgsl.jinja",
-              "bindings": "fc1ActivationScheduled",
               "dispatch": {
                 "x": "ceilDiv(interSize, decodeRows)",
                 "y": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * topK"
@@ -757,7 +377,18 @@
               "id": "output_stage",
               "name": "QMoE.OutputStageGemv",
               "shader": "qmoe-output-gemv.wgsl.jinja",
-              "bindings": "outputScheduled",
               "dispatch": {
                 "x": "ceilDiv(hiddenSize, decodeRows)",
                 "y": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens)"
@@ -781,8 +412,12 @@
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
-          "bindings": "route",
-          "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" }
         },
         {
           "id": "token_chunks",
@@ -792,20 +427,41 @@
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStage",
               "shader": "qmoe-fc1-activation-stage.wgsl.jinja",
-              "bindings": "fc1Activation",
               "dispatch": {
-                "threads": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * topK * interSize",
-                "workgroupSize": "constants.workgroupSize"
               }
             },
             {
               "id": "output_stage",
               "name": "QMoE.OutputStage",
               "shader": "qmoe-output-stage.wgsl.jinja",
-              "bindings": "output",
               "dispatch": {
-                "threads": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize",
-                "workgroupSize": "constants.workgroupSize"
               }
             }
           ]

   "domain": "com.microsoft",
   "name": "QMoE",
   "sinceVersion": 1,
+  "inputs": {
+    "inputT": { "onnx": "input", "dtype": "T" },
+    "routerT": { "onnx": "router_probs", "dtype": "T", "rank": 2 },
+    "fc1T": { "onnx": "fc1_experts_weights", "dtype": "T1", "rank": 3 },
+    "fc1ScalesT": { "onnx": "fc1_scales", "dtype": "T2" },
+    "fc2T": { "onnx": "fc2_experts_weights", "dtype": "T1", "rank": 3 },
+    "fc2ScalesT": { "onnx": "fc2_scales", "dtype": "T2" }
   },
+  "outputs": { "outputT": { "onnx": "output", "dtype": "T", "shape": "shapes.inputT" } },
+  "attributes": {
+    "activation_alpha": { "default": 1 },
+    "activation_beta": { "default": 0 },
+    "activation_type": { "default": "relu" },
+    "expert_weight_bits": { "default": 4 },
+    "k": { "default": 1 },
+    "normalize_routing_weights": { "default": 0 },
+    "quant_type": { "default": "int" },
+    "swiglu_fusion": { "default": 0 },
+    "use_sparse_mixer": { "default": 0 },
+    "weights_prepacked": { "default": -1 },
+    "block_size": {},
+    "swiglu_limit": {}
   },
   "attributeConstraints": {
     "activation_type": { "values": ["relu", "swiglu"] },
     "weights_prepacked": { "values": [-1, 0] }
   },
   "typeConstraints": { "T": ["float32"], "T1": ["uint8"], "T2": ["float32"] },
   "tunables": {
+    "workgroupSize": { "default": 64 },
+    "decodeLanes": { "default": 32 },
+    "decodeBlockTarget": { "default": 1024 },
+    "decodeMinLaneTrips": { "default": 4 },
+    "groupThreads": { "default": 8 },
+    "groupRegM": { "default": 4 },
+    "groupRegN": { "default": 4 },
+    "groupTileK": { "default": 16 },
+    "groupRouteWorkgroup": { "default": 256 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
     "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
     "wave32Effective": "wave32Adapter or pinSubgroupSize32",
+    "hiddenSize": "dim(shapes.inputT, ranks.inputT - 1)",
+    "numTokens": "numel(shapes.inputT) / max(1, hiddenSize)",
     "topK": "attrs.k",
     "weightBits": "attrs.expert_weight_bits",
     "quantBlockSize": "attrs.block_size if has(attrs, \"block_size\") else 0",
     "packSize": "2 if weightBits == 4 else 1",
     "quantMidpoint": "8 if weightBits == 4 else 128",
     "fusionSize": "2 if attrs.activation_type == \"swiglu\" and attrs.swiglu_fusion == 1 else 1",
+    "interSize": "dim(shapes.fc1T, 1) / fusionSize",
+    "fc1PackedCols": "dim(shapes.fc1T, 2)",
+    "fc2PackedCols": "dim(shapes.fc2T, 2)",
     "colWiseScales": "quantBlockSize == 0",
     "fc1ScaleBlocks": "1 if colWiseScales else hiddenSize / max(1, quantBlockSize)",
     "fc2ScaleBlocks": "1 if colWiseScales else interSize / max(1, quantBlockSize)",
     "activationSupported": "(attrs.activation_type == \"relu\" and attrs.swiglu_fusion == 0) or (attrs.activation_type == \"swiglu\" and attrs.swiglu_fusion == 1)",
     "routingModeSupported": "attrs.normalize_routing_weights == 0 or attrs.normalize_routing_weights == 1",
     "rawWeightLayout": "attrs.weights_prepacked == -1 or attrs.weights_prepacked == 0",
+    "inputOutputShapeOk": "((ranks.inputT == 2 and ranks.outputT == 2 and dim(shapes.outputT, 0) == dim(shapes.inputT, 0) and dim(shapes.outputT, 1) == dim(shapes.inputT, 1)) or (ranks.inputT == 3 and ranks.outputT == 3 and dim(shapes.outputT, 0) == dim(shapes.inputT, 0) and dim(shapes.outputT, 1) == dim(shapes.inputT, 1) and dim(shapes.outputT, 2) == dim(shapes.inputT, 2))) and hiddenSize > 0",
     "quantBlockSizeOk": "colWiseScales or (quantBlockSize >= 16 and pow2ceil(quantBlockSize) == quantBlockSize and hiddenSize % quantBlockSize == 0 and interSize % quantBlockSize == 0)",
+    "quantScalesOk": "tensorDtypes.fc1ScalesT == \"float32\" and tensorDtypes.fc2ScalesT == \"float32\" and dim(shapes.fc1ScalesT, 0) == dim(shapes.routerT, 1) and dim(shapes.fc2ScalesT, 0) == dim(shapes.routerT, 1) and dim(shapes.fc1ScalesT, 1) == dim(shapes.fc1T, 1) and dim(shapes.fc2ScalesT, 1) == hiddenSize and ((ranks.fc1ScalesT == 2 and ranks.fc2ScalesT == 2) if colWiseScales else (ranks.fc1ScalesT == 3 and ranks.fc2ScalesT == 3 and dim(shapes.fc1ScalesT, 2) == fc1ScaleBlocks and dim(shapes.fc2ScalesT, 2) == fc2ScaleBlocks))",
+    "quantShapeOk": "inputOutputShapeOk and ranks.routerT == 2 and ranks.fc1T == 3 and ranks.fc2T == 3 and dim(shapes.routerT, 0) == numTokens and dim(shapes.fc1T, 0) == dim(shapes.routerT, 1) and dim(shapes.fc2T, 0) == dim(shapes.routerT, 1) and dim(shapes.fc1T, 1) % fusionSize == 0 and dim(shapes.fc2T, 1) == hiddenSize and dim(shapes.fc1T, 2) * packSize == hiddenSize and dim(shapes.fc2T, 2) * packSize == interSize and quantBlockSizeOk and quantScalesOk",
+    "quantContract": "activationSupported and routingModeSupported and rawWeightLayout and quantShapeOk and topK >= 1 and topK <= dim(shapes.routerT, 1)",
     "workgroupSizeOk": "tunables.workgroupSize >= 1 and tunables.workgroupSize <= deviceWorkgroupCap",
     "hiddenChunkFits": "topK * interSize * 4 <= device.limits.maxStorageBufferBindingSize and topK * interSize * 4 <= device.limits.maxBufferSize",
     "hiddenChunkTokens": "numTokens if interSize == 0 else min(numTokens, max(1, floor(min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) / (topK * interSize * 4))))",
     "groupTileKVec": "ceilDiv(tunables.groupTileK, 4)",
     "groupThreadCount": "tunables.groupThreads * tunables.groupThreads",
     "groupSlots": "hiddenChunkTokens * topK",
+    "groupMaxTiles": "ceilDiv(groupSlots, max(1, groupTileM)) + dim(shapes.routerT, 1)",
     "groupSharedBytes": "(groupTileM * tunables.groupTileK + 2 * groupTileN * tunables.groupTileK + groupTileM) * 4",
     "groupSlotOutBytes": "hiddenChunkTokens * topK * hiddenSize * 4",
+    "groupedDeviceOk": "groupThreadCount <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.groupThreads <= device.limits.maxComputeWorkgroupSizeX and tunables.groupThreads <= device.limits.maxComputeWorkgroupSizeY and tunables.groupRouteWorkgroup <= deviceWorkgroupCap and groupSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.routerT, 1) * 8 <= device.limits.maxComputeWorkgroupStorageSize and tunables.groupTileK % 4 == 0",
+    "groupedDispatchOk": "groupMaxTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(interSize, max(1, groupTileN)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(hiddenSize, max(1, groupTileN)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "groupSlotOutFits": "groupSlotOutBytes <= device.limits.maxStorageBufferBindingSize and groupSlotOutBytes <= device.limits.maxBufferSize",
+    "groupedShapeOk": "numTokens * topK * 4 >= groupTileM * dim(shapes.routerT, 1)",
     "groupedContract": "quantContract and workgroupSizeOk and hiddenChunkFits and routedScratchFits and interSize > 0 and groupedDeviceOk and groupedDispatchOk and groupSlotOutFits and groupedShapeOk",
     "sgmatWorkgroup": "128",
     "sgmatSubgroups": "4",
     "sgmatRowSubtiles": "2",
+    "groupedSgmatStageElements": "max(groupTileM * 32, sgmatSubgroups * 4 * 64)",
+    "groupedSgmatSharedBytes": "(groupedSgmatStageElements + 64 * 32) * 4",
+    "groupedSgmatOk": "groupedContract and groupTileM == 32 and interSize % 32 == 0 and groupedSgmatSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(hiddenSize, 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and hiddenSize % 32 == 0",
     "decodeLanesOk": "tunables.decodeLanes >= 1 and tunables.decodeLanes <= tunables.workgroupSize and tunables.workgroupSize % tunables.decodeLanes == 0",
     "decodeRows": "max(1, tunables.workgroupSize / max(1, tunables.decodeLanes))",
     "decodeDeviceOk": "workgroupSizeOk and decodeLanesOk and tunables.decodeLanes <= device.limits.maxComputeWorkgroupSizeX and decodeRows <= device.limits.maxComputeWorkgroupSizeY and tunables.workgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
+    "decodeDispatchOk": "ceilDiv(interSize, decodeRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(hiddenSize, decodeRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and hiddenChunkTokens * topK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "decodeLaneDepth": "min(hiddenSize, interSize) / max(1, tunables.decodeLanes)",
     "decodeDepthOk": "decodeLaneDepth >= tunables.decodeMinLaneTrips",
     "splitFc1Blocks": "ceilDiv(numTokens * topK * interSize, max(1, tunables.workgroupSize))",
     "decodeOccupancyOk": "splitFc1Blocks <= tunables.decodeBlockTarget",
+    "decodeContract": "quantContract and decodeDeviceOk and decodeDispatchOk and hiddenChunkFits and routedScratchFits and interSize > 0 and decodeDepthOk and (decodeOccupancyOk or not groupedContract)",
     "tokens": "numTokens",
     "hidden": "hiddenSize",
+    "experts": "dim(shapes.routerT, 1)",
+    "fc1Rows": "dim(shapes.fc1T, 1)",
     "inter": "interSize",
     "workgroupSize": "tunables.workgroupSize",
     "activationType": "attrs.activation_type",
     "hasSwigluLimit": "has(attrs, \"swiglu_limit\")",
+    "swigluLimit": "attrs.swiglu_limit if has(attrs, \"swiglu_limit\") else 0"
   },
+  "bindings": {
+    "output": { "arg": "outputT", "buffer": "storage", "elementType": "f32" },
+    "router_probs": { "arg": "routerT", "buffer": "read-only-storage", "elementType": "f32" },
+    "route_expert": { "scratch": "routeExpert", "buffer": "storage", "elementType": "u32" },
+    "route_mix": { "scratch": "routeMix", "buffer": "storage", "elementType": "f32" },
+    "route_expert_2": {
+      "scratch": "routeExpert",
+      "name": "route_expert",
+      "buffer": "read-only-storage",
+      "elementType": "u32"
+    },
+    "slot_list": { "scratch": "slotList", "buffer": "storage", "elementType": "u32" },
+    "tile_meta": { "scratch": "tileMeta", "buffer": "storage", "elementType": "u32" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" },
+        {
+          "name": "tokenCount",
+          "type": "u32",
+          "value": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens)"
         }
+      ]
+    },
+    "input": { "arg": "inputT", "buffer": "read-only-storage", "elementType": "f32" },
+    "slot_list_2": { "scratch": "slotList", "name": "slot_list", "buffer": "read-only-storage", "elementType": "u32" },
+    "tile_meta_2": { "scratch": "tileMeta", "name": "tile_meta", "buffer": "read-only-storage", "elementType": "u32" },
+    "fc1_experts_weights": { "arg": "fc1T", "buffer": "read-only-storage", "elementType": "u32" },
+    "fc1_scales": { "arg": "fc1ScalesT", "buffer": "read-only-storage", "elementType": "f32" },
+    "hidden_act": { "scratch": "hiddenAct", "buffer": "storage", "elementType": "f32" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" },
+        { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" },
+        { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }
+      ]
+    },
+    "hidden_act_2": {
+      "scratch": "hiddenAct",
+      "name": "hidden_act",
+      "buffer": "read-only-storage",
+      "elementType": "f32"
+    },
+    "fc2_experts_weights": { "arg": "fc2T", "buffer": "read-only-storage", "elementType": "u32" },
+    "fc2_scales": { "arg": "fc2ScalesT", "buffer": "read-only-storage", "elementType": "f32" },
+    "slot_out": { "scratch": "slotOut", "buffer": "storage", "elementType": "f32" },
+    "slot_out_2": { "scratch": "slotOut", "name": "slot_out", "buffer": "read-only-storage", "elementType": "f32" },
+    "route_mix_2": { "scratch": "routeMix", "name": "route_mix", "buffer": "read-only-storage", "elementType": "f32" }
   },
   "variants": [
     {
         {
           "id": "output_stage",
           "name": "QMoE.OutputStageZeroInter",
+          "shader": "qmoe-output-zero-inter.wgsl.jinja",
+          "derive": { "outputElementCount": "numel(shapes.outputT)" },
+          "bindings": ["output"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.outputT)), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.outputT)), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "id": "quant_grouped_sgmat_routed",
       "priority": 32,
       "when": ["groupedSgmatOk", "wave32Effective"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
+      "derive": { "groupRouteWorkgroup": "tunables.groupRouteWorkgroup" },
       "intermediates": [
         { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" },
         { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" },
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": ["router_probs", "route_expert", "route_mix"],
+          "dispatch": {
+            "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
+            "z": 1
+          }
         },
         {
           "id": "token_chunks",
               "id": "group_stage",
               "name": "QMoE.GroupStage",
               "shader": "expert-group-slots.wgsl.jinja",
+              "bindings": ["route_expert_2", "slot_list", "tile_meta", "params"],
               "dispatch": {
+                "x": "min(ceilDiv((groupRouteWorkgroup), (groupRouteWorkgroup)), 65535)",
+                "y": "ceilDiv(ceilDiv((groupRouteWorkgroup), (groupRouteWorkgroup)), 65535)",
+                "z": 1
               }
             },
             {
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGroupedSgmat",
               "shader": "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja",
+              "bindings": ["input", "slot_list_2", "tile_meta_2", "fc1_experts_weights", "fc1_scales", "hidden_act", "params_2"],
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize * fusionSize, 64)" }
             },
             {
               "id": "output_stage",
               "name": "QMoE.OutputStageGroupedSgmat",
               "shader": "qmoe-output-grouped-sgmat.wgsl.jinja",
+              "bindings": ["hidden_act_2", "slot_list_2", "tile_meta_2", "fc2_experts_weights", "fc2_scales", "slot_out"],
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenSize, 64)" }
             },
             {
               "id": "mix_stage",
               "name": "QMoE.MixStage",
               "shader": "qmoe-mix-stage.wgsl.jinja",
+              "bindings": ["slot_out_2", "route_mix_2", "output", "params"],
               "dispatch": {
+                "x": "min(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
+                "y": "ceilDiv(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
+                "z": 1
               }
             }
           ]
         }
+      ]
     },
     {
       "id": "quant_grouped_routed",
       "priority": 30,
       "when": ["groupedContract"],
+      "derive": {
         "groupTileK": "tunables.groupTileK",
         "groupThreads": "tunables.groupThreads",
         "regM": "tunables.groupRegM",
         "regN": "tunables.groupRegN",
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": ["router_probs", "route_expert", "route_mix"],
+          "dispatch": {
+            "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
+            "z": 1
+          }
         },
         {
           "id": "token_chunks",
               "id": "group_stage",
               "name": "QMoE.GroupStage",
               "shader": "expert-group-slots.wgsl.jinja",
+              "bindings": ["route_expert_2", "slot_list", "tile_meta", "params"],
               "dispatch": {
+                "x": "min(ceilDiv((groupRouteWorkgroup), (groupRouteWorkgroup)), 65535)",
+                "y": "ceilDiv(ceilDiv((groupRouteWorkgroup), (groupRouteWorkgroup)), 65535)",
+                "z": 1
               }
             },
             {
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGrouped",
               "shader": "qmoe-fc1-activation-grouped.wgsl.jinja",
+              "bindings": ["input", "slot_list_2", "tile_meta_2", "fc1_experts_weights", "fc1_scales", "hidden_act", "params_2"],
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" }
             },
             {
               "id": "output_stage",
               "name": "QMoE.OutputStageGrouped",
               "shader": "qmoe-output-grouped.wgsl.jinja",
+              "bindings": ["hidden_act_2", "slot_list_2", "tile_meta_2", "fc2_experts_weights", "fc2_scales", "slot_out"],
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenSize, groupTileN)" }
             },
             {
               "id": "mix_stage",
               "name": "QMoE.MixStage",
               "shader": "qmoe-mix-stage.wgsl.jinja",
+              "bindings": ["slot_out_2", "route_mix_2", "output", "params"],
               "dispatch": {
+                "x": "min(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
+                "y": "ceilDiv(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
+                "z": 1
               }
             }
           ]
       "id": "quant_gemv_routed",
       "priority": 20,
       "when": ["decodeContract"],
+      "derive": { "decodeLanes": "tunables.decodeLanes" },
       "intermediates": [
         { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" },
         { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" },
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": ["router_probs", "route_expert", "route_mix"],
+          "dispatch": {
+            "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
+            "z": 1
+          }
         },
         {
           "id": "token_chunks",
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGemv",
               "shader": "qmoe-fc1-activation-gemv.wgsl.jinja",
+              "bindings": ["input", "route_expert_2", "fc1_experts_weights", "fc1_scales", "hidden_act", "params_2"],
               "dispatch": {
                 "x": "ceilDiv(interSize, decodeRows)",
                 "y": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * topK"
               "id": "output_stage",
               "name": "QMoE.OutputStageGemv",
               "shader": "qmoe-output-gemv.wgsl.jinja",
+              "bindings": [
+                "hidden_act_2",
+                "route_expert_2",
+                "route_mix_2",
+                "fc2_experts_weights",
+                "fc2_scales",
+                "output",
+                {
+                  "name": "params",
+                  "struct": [{ "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }]
+                }
+              ],
               "dispatch": {
                 "x": "ceilDiv(hiddenSize, decodeRows)",
                 "y": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens)"
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": ["router_probs", "route_expert", "route_mix"],
+          "dispatch": {
+            "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
+            "z": 1
+          }
         },
         {
           "id": "token_chunks",
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStage",
               "shader": "qmoe-fc1-activation-stage.wgsl.jinja",
+              "bindings": [
+                "input",
+                "route_expert_2",
+                "fc1_experts_weights",
+                "fc1_scales",
+                "hidden_act",
+                {
+                  "name": "params",
+                  "struct": [
+                    { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" },
+                    { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" },
+                    { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" },
+                    {
+                      "name": "tokenCount",
+                      "type": "u32",
+                      "value": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens)"
+                    }
+                  ]
+                }
+              ],
               "dispatch": {
+                "x": "min(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * topK * interSize), (workgroupSize)), 65535)",
+                "y": "ceilDiv(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * topK * interSize), (workgroupSize)), 65535)",
+                "z": 1
               }
             },
             {
               "id": "output_stage",
               "name": "QMoE.OutputStage",
               "shader": "qmoe-output-stage.wgsl.jinja",
+              "bindings": ["hidden_act_2", "route_expert_2", "route_mix_2", "fc2_experts_weights", "fc2_scales", "output", "params"],
               "dispatch": {
+                "x": "min(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
+                "y": "ceilDiv(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
+                "z": 1
               }
             }
           ]

build/webgpu/metadata.json CHANGED Viewed

@@ -1,29 +1,38 @@
 {
   "name": "com.microsoft.QMoE",
-  "id": "_com_microsoft_qmoe_webgpu_6caed83",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "zBVupnayA0Ggc4y737P9tzqyM/q/N8F4j+FEJXz6Wks=",
       "expert-group-slots.wgsl.jinja": "Ta+3H2FA1qRRzksdkwgKLZiO+AokmuM+B9JVN8cv1EQ=",
-      "manifest.json": "F9s2iIz0nNI2vJ+2zAFr55u8jaKjAEe5qvfmGxSc6mE=",
       "qmoe-fc1-activation-gemv.wgsl.jinja": "UvP8Vyx1E9+/+I09w/z9iB6h6MHrSrOmjKfIelyIS0M=",
-      "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja": "r/vlFM8WL1WfjkX7P2/DjKxrMDRKBmI+PtV/qQVLgCc=",
-      "qmoe-fc1-activation-grouped.wgsl.jinja": "y2vGeNeXE5YmNrM2WrXogyF/fTvQ/nqYR4uWBTDvpPY=",
-      "qmoe-fc1-activation-stage.wgsl.jinja": "Hx7fYjr8CwB2jXeTnh6G6IyavWrLwyOdtUsJwVEpgGk=",
-      "qmoe-mix-stage.wgsl.jinja": "+deuFyuBn04UeIX8Hk2/ntCHNfM1vgkNDTcsh3uXVdI=",
       "qmoe-output-gemv.wgsl.jinja": "OJLoDDWMYE5yaTeDGY7qUeuf6AepJL7ZZLx8ko6OTEQ=",
-      "qmoe-output-grouped-sgmat.wgsl.jinja": "aoBOyGPOM2MXOpjvVCppacvWra4bVBewQEbjL2MPe2k=",
-      "qmoe-output-grouped.wgsl.jinja": "mHkrsubKwd0YWYo4TMpl4HWdbWheoeiAsUU2WsUON8E=",
-      "qmoe-output-stage.wgsl.jinja": "f9CWE9XR/5RTULNXdTHwA4gFsOL1QRg0H/rT3dD32Gs=",
-      "qmoe-output-zero-inter.wgsl.jinja": "6iBjk5YqxSSoXufDo/aJJRS4v8ZOp7wF2aRRXU4m1SA=",
-      "qmoe-route-stage.wgsl.jinja": "cGnJJB0bEdlcPFOKm8s42byI9ePixYfmfK9nK05L7V4=",
-      "test.json": "3i3FSregL3rNfwD/Wi7X2Ix296NNB1zIV+Ap9TrpMeU="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.QMoE" }
 }

 {
   "name": "com.microsoft.QMoE",
+  "id": "_com_microsoft_qmoe_webgpu_b3768d9",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "YuXSEWa4U+6Ek0E0ZkIeDrzqQ672IelokTilKJuKHf0=",
       "expert-group-slots.wgsl.jinja": "Ta+3H2FA1qRRzksdkwgKLZiO+AokmuM+B9JVN8cv1EQ=",
+      "manifest.json": "r3eaOyzncaumYYrgh5XHHfvd9mHdvVPR3cl5xbiwDDA=",
       "qmoe-fc1-activation-gemv.wgsl.jinja": "UvP8Vyx1E9+/+I09w/z9iB6h6MHrSrOmjKfIelyIS0M=",
+      "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja": "uWNSXPm0htntxHKCNBzDIo5qFmROP9wDF3mjBGQo+Wo=",
+      "qmoe-fc1-activation-grouped.wgsl.jinja": "6NEYMHvSNgCp6KHJveI4LELrTqr8u/lxIJjSkRX68V4=",
+      "qmoe-fc1-activation-stage.wgsl.jinja": "C9DvNGQ2u4eBAAceyEdbGX13goQFJbxSunkXXhYN604=",
+      "qmoe-mix-stage.wgsl.jinja": "+xEc6boXboFbx+4AJ2/KWElM04LRuufmiIKfEK51/oo=",
       "qmoe-output-gemv.wgsl.jinja": "OJLoDDWMYE5yaTeDGY7qUeuf6AepJL7ZZLx8ko6OTEQ=",
+      "qmoe-output-grouped-sgmat.wgsl.jinja": "jopTCpUPoIw9MBRhrmjQuZioTbJJIygwBvbwRyhHvUY=",
+      "qmoe-output-grouped.wgsl.jinja": "dOoFq8c0+yBCwOae+QVKBQyodJUK4vwam1RQb46oBNU=",
+      "qmoe-output-stage.wgsl.jinja": "yuSg3G05eKQOunyR5lcCEaJ5y5lXoT73RzT9rq6ahJQ=",
+      "qmoe-output-zero-inter.wgsl.jinja": "6jHYpVxU5+9WMflhwMste4v17hq7WWA1BDDS4q4DLS4=",
+      "qmoe-route-stage.wgsl.jinja": "SF5VN6d0zp/ssCGnTwCWccgbnOkZ0mptZiwIbMX3gXA=",
+      "test.json": "/ffqKBtjk2jseUZPswe8gOyYZfhNJu9aOZg8GNlB+9A="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "quant_zero_inter": ["qmoe-output-zero-inter.wgsl.jinja"],
+      "quant_grouped_sgmat_routed": ["expert-group-slots.wgsl.jinja", "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja", "qmoe-mix-stage.wgsl.jinja", "qmoe-output-grouped-sgmat.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"],
+      "quant_grouped_routed": ["expert-group-slots.wgsl.jinja", "qmoe-fc1-activation-grouped.wgsl.jinja", "qmoe-mix-stage.wgsl.jinja", "qmoe-output-grouped.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"],
+      "quant_gemv_routed": ["qmoe-fc1-activation-gemv.wgsl.jinja", "qmoe-output-gemv.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"],
+      "quant_split_routed": ["qmoe-fc1-activation-stage.wgsl.jinja", "qmoe-output-stage.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/qmoe-fc1-activation-grouped-sgmat.wgsl.jinja CHANGED Viewed

@@ -5,14 +5,13 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
-// com.microsoft.QMoE, expert-grouped FC1 + activation stage on the subgroup-matrix tier.
-//
-// Same decomposition as the scalar grouped stage -- routed slots arrive sorted by expert,
-// so a tile shares one expert and its weight tile is dequantized once and reused down the
-// tile -- with the inner product on the matrix units and the packed weights dequantized
-// straight into the shared B operand.
 //
 // A fused SwiGLU takes its gate from weight row 2c and its up from row 2c+1. Those are
 // staged as adjacent columns of one B tile rather than as two tiles. The result columns a
@@ -21,7 +20,6 @@ const HIDDEN: u32 = {{ hidden }}u;
 const FC1_ROWS: u32 = {{ fc1Rows }}u;
 const INTER: u32 = {{ inter }}u;
 const TOP_K: u32 = {{ topK }}u;
-const TILE_ROWS: u32 = {{ groupTileM }}u;
 const TILE_COLS: u32 = 64u;
 const TILE_K: u32 = 32u;
 const SUB_COLS: u32 = 32u;
@@ -43,7 +41,7 @@ const FC1_BLOCKS: u32 = {{ fc1ScaleBlocks }}u;
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
-// the caller guarantees `col` is a multiple of four.
 fn fc1_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
   let base = expert * FC1_ROWS + row;
 {% if packSize == 2 %}
@@ -86,16 +84,16 @@ fn activate(value: f32) -> f32 {
 {% endif %}
-var<workgroup> row_slot: array<u32, {{ groupTileM }}>;
-var<workgroup> tile_A: array<f32, {{ groupTileM }} * 32>;
 var<workgroup> tile_B: array<f32, 64 * 32>;
-var<workgroup> scratch: array<array<array<f32, 64>, 4>, {{ sgmatSubgroups }}>;
 // Rows past the tile's end stage zero and are dropped at the store, so a short final
 // tile needs no separate arm.
-fn loadSHMA(k_idx: u32, rows: u32, row: u32, c_idx: u32) {
   let col = c_idx * 8u;
-  let base = (params.tokenOffset + row_slot[row] / TOP_K) * HIDDEN + k_idx + col;
   if (row < rows) {
     for (var i = 0u; i < 8u; i = i + 1u) {
       tile_A[row * TILE_K + col + i] = f32(input[base + i]);
@@ -143,10 +141,7 @@ fn main(
   let rows = tile_meta[1u + tile * 3u + 2u];
   let weight_base = wid.y * TILE_COLS;
-  for (var m = local_idx; m < TILE_ROWS; m = m + {{ sgmatWorkgroup }}u) {
-    row_slot[m] = select(0u, slot_list[slice_base + m], m < rows);
-  }
-  workgroupBarrier();
   let subtile_id = local_idx / sg_size;
   let subtile_idx = subtile_id / {{ sgmatRowSubtiles }}u;
@@ -159,20 +154,21 @@ fn main(
 {% endfor %}{% endfor %}
   for (var kidx = 0u; kidx < HIDDEN; kidx = kidx + TILE_K) {
-    loadSHMA(kidx, rows, local_idx / 4u, local_idx % 4u);
     loadSHMB(expert, weight_base, kidx, local_idx / 2u, local_idx % 2u);
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
-{% set operandScalar = "f32" %}      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
-      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
-      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
-      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
-      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
-      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
-      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -197,20 +193,20 @@ fn main(
   // The scratch banks are shared by both row groups, so the readback is fenced on both
   // sides of each collective store.
 {% for c in [0, 1, 2, 3] %}
-  subgroupMatrixStore<row_major>(&scratch[subtile_id][{{ c }}], 0u, matC{{ group }}{{ c }}, 8u);
 {% endfor %}
   workgroupBarrier();
   {
     let m = base_A + {{ 8 * group }}u + row;
     if (m < rows) {
-      let out_base = row_slot[m] * INTER;
 {% for c in [0, 1, 2, 3] %}
 {% if isSwiglu %}
       {
         let out_col = (out_weight_base + col + {{ 8 * c }}u) / 2u;
         if (out_col < INTER) {
-          hidden_act[out_base + out_col] = activate(scratch[subtile_id][{{ c }}][row * 8u + col],
-                                                    scratch[subtile_id][{{ c }}][row * 8u + col2]);
         }
       }
 {% else %}
@@ -218,7 +214,7 @@ fn main(
       {
         let out_col = out_weight_base + {{ half }} + {{ 8 * c }}u;
         if (out_col < INTER) {
-          hidden_act[out_base + out_col] = activate(scratch[subtile_id][{{ c }}][row * 8u + {{ half }}]);
         }
       }
 {% endfor %}

 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
+// Expert-grouped FC1 and activation stage using subgroup-matrix operations.
+// Routed slots arrive sorted by expert, so each tile shares one expert and
+// reuses its dequantized weight tile across rows. Packed weights are
+// dequantized directly into the shared B operand.
 //
 // A fused SwiGLU takes its gate from weight row 2c and its up from row 2c+1. Those are
 // staged as adjacent columns of one B tile rather than as two tiles. The result columns a
 const FC1_ROWS: u32 = {{ fc1Rows }}u;
 const INTER: u32 = {{ inter }}u;
 const TOP_K: u32 = {{ topK }}u;
 const TILE_COLS: u32 = 64u;
 const TILE_K: u32 = 32u;
 const SUB_COLS: u32 = 32u;
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
+// this helper requires `col` to be a multiple of four.
 fn fc1_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
   let base = expert * FC1_ROWS + row;
 {% if packSize == 2 %}
 {% endif %}
+// The final K-loop barrier ends the A operand's lifetime. Its storage then
+// holds four 8x8 output banks per subgroup, fenced by the epilogue barriers.
+var<workgroup> tile_A: array<f32, {{ groupedSgmatStageElements }}>;
 var<workgroup> tile_B: array<f32, 64 * 32>;
 // Rows past the tile's end stage zero and are dropped at the store, so a short final
 // tile needs no separate arm.
+fn loadSHMA(k_idx: u32, rows: u32, row: u32, c_idx: u32, a_slot: u32) {
   let col = c_idx * 8u;
+  let base = (params.tokenOffset + a_slot / TOP_K) * HIDDEN + k_idx + col;
   if (row < rows) {
     for (var i = 0u; i < 8u; i = i + 1u) {
       tile_A[row * TILE_K + col + i] = f32(input[base + i]);
   let rows = tile_meta[1u + tile * 3u + 2u];
   let weight_base = wid.y * TILE_COLS;
+  let a_slot = select(0u, slot_list[slice_base + local_idx / 4u], local_idx / 4u < rows);
   let subtile_id = local_idx / sg_size;
   let subtile_idx = subtile_id / {{ sgmatRowSubtiles }}u;
 {% endfor %}{% endfor %}
   for (var kidx = 0u; kidx < HIDDEN; kidx = kidx + TILE_K) {
+    loadSHMA(kidx, rows, local_idx / 4u, local_idx % 4u, a_slot);
     loadSHMB(expert, weight_base, kidx, local_idx / 2u, local_idx % 2u);
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set operandScalar = "f32" %}{% set directInputs = directMatrixInputs is defined and directMatrixInputs %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+{% for r in range(2) %}
+      var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});
+{% endfor %}
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+{% for c in range(4) %}
+      var matB{{ c }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&{{ "xm" if directInputs else "tile_B" }}, matrix_b_offset{% if c > 0 %} + {{ c * 8 }}u * TILE_K{% endif %}, {{ "N" if directInputs else "TILE_K" }});
+{% endfor %}
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
   // The scratch banks are shared by both row groups, so the readback is fenced on both
   // sides of each collective store.
 {% for c in [0, 1, 2, 3] %}
+  subgroupMatrixStore<row_major>(&tile_A, (subtile_id * 4u + {{ c }}u) * 64u, matC{{ group }}{{ c }}, 8u);
 {% endfor %}
   workgroupBarrier();
   {
     let m = base_A + {{ 8 * group }}u + row;
     if (m < rows) {
+      let out_base = slot_list[slice_base + m] * INTER;
 {% for c in [0, 1, 2, 3] %}
 {% if isSwiglu %}
       {
         let out_col = (out_weight_base + col + {{ 8 * c }}u) / 2u;
         if (out_col < INTER) {
+          hidden_act[out_base + out_col] = activate(tile_A[(subtile_id * 4u + {{ c }}u) * 64u + row * 8u + col],
+                                                    tile_A[(subtile_id * 4u + {{ c }}u) * 64u + row * 8u + col2]);
         }
       }
 {% else %}
       {
         let out_col = out_weight_base + {{ half }} + {{ 8 * c }}u;
         if (out_col < INTER) {
+          hidden_act[out_base + out_col] = activate(tile_A[(subtile_id * 4u + {{ c }}u) * 64u + row * 8u + {{ half }}]);
         }
       }
 {% endfor %}

build/webgpu/qmoe-fc1-activation-grouped.wgsl.jinja CHANGED Viewed

@@ -43,7 +43,7 @@ fn fc1_weight(expert: u32, row: u32, col: u32) -> f32 {
 }
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
-// the caller guarantees `col` is a multiple of four.
 fn fc1_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
   let base = expert * FC1_ROWS + row;
 {% if packSize == 2 %}

 }
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
+// this helper requires `col` to be a multiple of four.
 fn fc1_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
   let base = expert * FC1_ROWS + row;
 {% if packSize == 2 %}

build/webgpu/qmoe-fc1-activation-stage.wgsl.jinja CHANGED Viewed

@@ -60,11 +60,10 @@ fn activate(value: f32) -> f32 {
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the maxComputeWorkgroupsPerDimension
-  // workgroup-per-dimension dispatch limit. Reduces to gid.x when nwg.y == 1.
-  let index = gid.x + gid.y * nwg.x * WG;
 {% if topK > 1 %}
   // One activation row per routed slot rather than per token: the grid is
   // (token, slot, inter) and hidden_act is addressed by the same flat index.

 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  // 2D-folded flat index: gid.y carries the high bits past the per-axis dispatch fold width.
+  // Reduces to gid.x when the dispatch does not fold.
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
 {% if topK > 1 %}
   // One activation row per routed slot rather than per token: the grid is
   // (token, slot, inter) and hidden_act is addressed by the same flat index.

build/webgpu/qmoe-mix-stage.wgsl.jinja CHANGED Viewed

@@ -10,11 +10,10 @@ const TOP_K: u32 = {{ topK }}u;
 const WG: u32 = {{ workgroupSize }}u;
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the maxComputeWorkgroupsPerDimension
-  // workgroup-per-dimension dispatch limit. Reduces to gid.x when nwg.y == 1.
-  let index = gid.x + gid.y * nwg.x * WG;
   let total = params.tokenCount * HIDDEN;
   if (index >= total) {
     return;

 const WG: u32 = {{ workgroupSize }}u;
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  // 2D-folded flat index: gid.y carries the high bits past the per-axis dispatch fold width.
+  // Reduces to gid.x when the dispatch does not fold.
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let total = params.tokenCount * HIDDEN;
   if (index >= total) {
     return;

build/webgpu/qmoe-output-grouped-sgmat.wgsl.jinja CHANGED Viewed

@@ -5,22 +5,18 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
-// com.microsoft.QMoE, expert-grouped FC2 stage on the subgroup-matrix tier.
-//
-// Same decomposition as the scalar grouped stage -- the group stage has sorted routed
-// slots by expert, so one tile shares an expert and its weight tile is dequantized once
-// and reused across the tile's rows -- but the inner product runs on the matrix units.
-// The packed weights are dequantized straight into the shared B tile, so no float weight
-// matrix is ever materialized, and the accumulate stays f32, which is exact for a 4-bit
-// code times its scale.
 //
 // The M axis is the tile's routed-slot slice rather than a contiguous row range, so the
 // A loader gathers through `slot_list` and the store scatters back through it.
 const HIDDEN: u32 = {{ hidden }}u;
 const INTER: u32 = {{ inter }}u;
-const TILE_ROWS: u32 = {{ groupTileM }}u;
 const TILE_COLS: u32 = 64u;
 const TILE_K: u32 = 32u;
 const SUB_COLS: u32 = 32u;
@@ -39,7 +35,7 @@ const FC2_BLOCKS: u32 = {{ fc2ScaleBlocks }}u;
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
-// the caller guarantees `col` is a multiple of four.
 fn fc2_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
   let base = expert * HIDDEN + row;
 {% if packSize == 2 %}
@@ -62,16 +58,16 @@ fn fc2_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
 }
-var<workgroup> row_slot: array<u32, {{ groupTileM }}>;
-var<workgroup> tile_A: array<f32, {{ groupTileM }} * 32>;
 var<workgroup> tile_B: array<f32, 64 * 32>;
-var<workgroup> scratch: array<array<array<f32, 64>, 4>, {{ sgmatSubgroups }}>;
 // Rows past the tile's end stage zero and are dropped at the store, so a short final
 // tile needs no separate arm.
-fn loadSHMA(k_idx: u32, rows: u32, row: u32, c_idx: u32) {
   let col = c_idx * 8u;
-  let base = row_slot[row] * INTER + k_idx + col;
   if (row < rows) {
     for (var i = 0u; i < 8u; i = i + 1u) {
       tile_A[row * TILE_K + col + i] = hidden_act[base + i];
@@ -105,15 +101,15 @@ fn loadSHMB(expert: u32, col_base: u32, k_idx: u32, row: u32, c_idx: u32) {
 fn storeCol(col_base: u32, row: u32, cc: u32, slot: u32, src_slot: u32, out_slot: u32) {
   let c = cc + slot * 8u;
   if (col_base + c < HIDDEN) {
-    slot_out[out_slot * HIDDEN + col_base + c] = scratch[src_slot][slot][row * 8u + cc];
   }
 }
-fn storeOutput(col_base: u32, row: u32, col: u32, src_slot: u32, row_base: u32, rows: u32) {
   let m = row_base + row;
   if (m < rows) {
     let col2 = col + 1u;
-    let out_slot = row_slot[m];
 {% for blockCol in [0, 1, 2, 3] %}
     storeCol(col_base, row, col,  {{ blockCol }}u, src_slot, out_slot);
     storeCol(col_base, row, col2, {{ blockCol }}u, src_slot, out_slot);
@@ -137,10 +133,7 @@ fn main(
   let rows = tile_meta[1u + tile * 3u + 2u];
   let col_base = wid.y * TILE_COLS;
-  for (var m = local_idx; m < TILE_ROWS; m = m + {{ sgmatWorkgroup }}u) {
-    row_slot[m] = select(0u, slot_list[slice_base + m], m < rows);
-  }
-  workgroupBarrier();
   let subtile_id = local_idx / sg_size;
   let subtile_idx = subtile_id / {{ sgmatRowSubtiles }}u;
@@ -153,20 +146,21 @@ fn main(
 {% endfor %}{% endfor %}
   for (var kidx = 0u; kidx < INTER; kidx = kidx + TILE_K) {
-    loadSHMA(kidx, rows, local_idx / 4u, local_idx % 4u);
     loadSHMB(expert, col_base, kidx, local_idx / 2u, local_idx % 2u);
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
-{% set operandScalar = "f32" %}      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
-      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
-      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
-      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
-      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
-      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
-      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -183,17 +177,17 @@ fn main(
   // Each scratch bank is written by a collective store and then read across lanes, and
   // the two row groups reuse the same banks, so both barriers are load-bearing.
 {% for col in [0, 1, 2, 3] %}
-  subgroupMatrixStore<row_major>(&scratch[subtile_id][{{ col }}], 0u, matC0{{ col }}, 8u);
 {% endfor %}
   workgroupBarrier();
   let row = sg_id / 4u;
   let col = (sg_id % 4u) * 2u;
-  storeOutput(col_base + base_B, row, col, subtile_id, base_A, rows);
   workgroupBarrier();
 {% for col in [0, 1, 2, 3] %}
-  subgroupMatrixStore<row_major>(&scratch[subtile_id][{{ col }}], 0u, matC1{{ col }}, 8u);
 {% endfor %}
   workgroupBarrier();
-  storeOutput(col_base + base_B, row, col, subtile_id, base_A + 8u, rows);
 }

 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
+// Expert-grouped FC2 stage using subgroup-matrix operations. Routed slots are
+// sorted by expert, so each tile shares one expert and reuses its dequantized
+// weight tile across rows. Packed weights are dequantized directly into the
+// shared B tile, and accumulation remains in f32.
 //
 // The M axis is the tile's routed-slot slice rather than a contiguous row range, so the
 // A loader gathers through `slot_list` and the store scatters back through it.
 const HIDDEN: u32 = {{ hidden }}u;
 const INTER: u32 = {{ inter }}u;
 const TILE_COLS: u32 = 64u;
 const TILE_K: u32 = 32u;
 const SUB_COLS: u32 = 32u;
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
+// this helper requires `col` to be a multiple of four.
 fn fc2_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
   let base = expert * HIDDEN + row;
 {% if packSize == 2 %}
 }
+// The final K-loop barrier ends the A operand's lifetime. Its storage then
+// holds four 8x8 output banks per subgroup, fenced by the epilogue barriers.
+var<workgroup> tile_A: array<f32, {{ groupedSgmatStageElements }}>;
 var<workgroup> tile_B: array<f32, 64 * 32>;
 // Rows past the tile's end stage zero and are dropped at the store, so a short final
 // tile needs no separate arm.
+fn loadSHMA(k_idx: u32, rows: u32, row: u32, c_idx: u32, a_slot: u32) {
   let col = c_idx * 8u;
+  let base = a_slot * INTER + k_idx + col;
   if (row < rows) {
     for (var i = 0u; i < 8u; i = i + 1u) {
       tile_A[row * TILE_K + col + i] = hidden_act[base + i];
 fn storeCol(col_base: u32, row: u32, cc: u32, slot: u32, src_slot: u32, out_slot: u32) {
   let c = cc + slot * 8u;
   if (col_base + c < HIDDEN) {
+    slot_out[out_slot * HIDDEN + col_base + c] = tile_A[(src_slot * 4u + slot) * 64u + row * 8u + cc];
   }
 }
+fn storeOutput(col_base: u32, row: u32, col: u32, src_slot: u32, row_base: u32, rows: u32, slice_base: u32) {
   let m = row_base + row;
   if (m < rows) {
     let col2 = col + 1u;
+    let out_slot = slot_list[slice_base + m];
 {% for blockCol in [0, 1, 2, 3] %}
     storeCol(col_base, row, col,  {{ blockCol }}u, src_slot, out_slot);
     storeCol(col_base, row, col2, {{ blockCol }}u, src_slot, out_slot);
   let rows = tile_meta[1u + tile * 3u + 2u];
   let col_base = wid.y * TILE_COLS;
+  let a_slot = select(0u, slot_list[slice_base + local_idx / 4u], local_idx / 4u < rows);
   let subtile_id = local_idx / sg_size;
   let subtile_idx = subtile_id / {{ sgmatRowSubtiles }}u;
 {% endfor %}{% endfor %}
   for (var kidx = 0u; kidx < INTER; kidx = kidx + TILE_K) {
+    loadSHMA(kidx, rows, local_idx / 4u, local_idx % 4u, a_slot);
     loadSHMB(expert, col_base, kidx, local_idx / 2u, local_idx % 2u);
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set operandScalar = "f32" %}{% set directInputs = directMatrixInputs is defined and directMatrixInputs %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+{% for r in range(2) %}
+      var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});
+{% endfor %}
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+{% for c in range(4) %}
+      var matB{{ c }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&{{ "xm" if directInputs else "tile_B" }}, matrix_b_offset{% if c > 0 %} + {{ c * 8 }}u * TILE_K{% endif %}, {{ "N" if directInputs else "TILE_K" }});
+{% endfor %}
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
   // Each scratch bank is written by a collective store and then read across lanes, and
   // the two row groups reuse the same banks, so both barriers are load-bearing.
 {% for col in [0, 1, 2, 3] %}
+  subgroupMatrixStore<row_major>(&tile_A, (subtile_id * 4u + {{ col }}u) * 64u, matC0{{ col }}, 8u);
 {% endfor %}
   workgroupBarrier();
   let row = sg_id / 4u;
   let col = (sg_id % 4u) * 2u;
+  storeOutput(col_base + base_B, row, col, subtile_id, base_A, rows, slice_base);
   workgroupBarrier();
 {% for col in [0, 1, 2, 3] %}
+  subgroupMatrixStore<row_major>(&tile_A, (subtile_id * 4u + {{ col }}u) * 64u, matC1{{ col }}, 8u);
 {% endfor %}
   workgroupBarrier();
+  storeOutput(col_base + base_B, row, col, subtile_id, base_A + 8u, rows, slice_base);
 }

build/webgpu/qmoe-output-grouped.wgsl.jinja CHANGED Viewed

@@ -41,7 +41,7 @@ fn fc2_weight(expert: u32, row: u32, col: u32) -> f32 {
 }
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
-// the caller guarantees `col` is a multiple of four.
 fn fc2_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
   let base = expert * HIDDEN + row;
 {% if packSize == 2 %}

 }
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
+// this helper requires `col` to be a multiple of four.
 fn fc2_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
   let base = expert * HIDDEN + row;
 {% if packSize == 2 %}

build/webgpu/qmoe-output-stage.wgsl.jinja CHANGED Viewed

@@ -35,11 +35,10 @@ fn fc2_weight(expert: u32, row: u32, col: u32) -> f32 {
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the maxComputeWorkgroupsPerDimension
-  // workgroup-per-dimension dispatch limit. Reduces to gid.x when nwg.y == 1.
-  let index = gid.x + gid.y * nwg.x * WG;
   let total = params.tokenCount * HIDDEN;
   if (index >= total) {
     return;

 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  // 2D-folded flat index: gid.y carries the high bits past the per-axis dispatch fold width.
+  // Reduces to gid.x when the dispatch does not fold.
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let total = params.tokenCount * HIDDEN;
   if (index >= total) {
     return;

build/webgpu/qmoe-output-zero-inter.wgsl.jinja CHANGED Viewed

@@ -3,10 +3,9 @@
 const WG: u32 = {{ workgroupSize }}u;
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  let index = gid.x + gid.y * nwg.x * WG;
-  if (index < {{ source.outputElementCount }}u) {
     output[index] = 0.0;
   }
 }

 const WG: u32 = {{ workgroupSize }}u;
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
+  if (index < {{ outputElementCount }}u) {
     output[index] = 0.0;
   }
 }

build/webgpu/qmoe-route-stage.wgsl.jinja CHANGED Viewed

@@ -79,11 +79,10 @@ fn top1_route(token: u32) -> Top1Route {
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the maxComputeWorkgroupsPerDimension
-  // workgroup-per-dimension dispatch limit. Reduces to gid.x when nwg.y == 1.
-  let token = gid.x + gid.y * nwg.x * WG;
   if (token >= TOKENS) {
     return;
   }

 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  // 2D-folded flat index: gid.y carries the high bits past the per-axis dispatch fold width.
+  // Reduces to gid.x when the dispatch does not fold.
+  let token = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   if (token >= TOKENS) {
     return;
   }

build/webgpu/test.json CHANGED Viewed

@@ -1,10 +1,9 @@
 {
-  "op": "com.microsoft.QMoE",
   "cases": [
     {
       "name": "q4_colwise_default_relu",
       "provenance": {
-        "notes": "Exercises every exact default at once: 4-bit integer weights, k=1, ReLU, column-wise scales, dense routing, and raw provider layout."
       },
       "inputs": {
         "inputT": {
@@ -39,7 +38,7 @@
     {
       "name": "q4_top1_tie_selects_higher_expert",
       "provenance": {
-        "notes": "Pins ONNX Runtime's descending `(logit, expert_index)` ordering with selected-weight normalization: equal logits select expert 1, whose hand-computed output is [16, 32], rather than the all-zero expert 0."
       },
       "attrs": { "normalize_routing_weights": 1 },
       "inputs": {
@@ -227,28 +226,16 @@
     {
       "name": "q4_zero_inter",
       "provenance": {
-        "notes": "A zero intermediate dimension has an exact zero result without binding empty weight buffers in a shader."
       },
       "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1 },
       "inputs": {
-        "inputT": {
-          "dtype": "float32",
-          "shape": [3, 4],
-          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.3 }
-        },
-        "routerT": {
-          "dtype": "float32",
-          "shape": [3, 3],
-          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
-        },
         "fc1T": { "dtype": "uint8", "shape": [3, 0, 2], "data": { "kind": "values", "values": [] } },
         "fc1ScalesT": { "dtype": "float32", "shape": [3, 0], "data": { "kind": "values", "values": [] } },
         "fc2T": { "dtype": "uint8", "shape": [3, 4, 0], "data": { "kind": "values", "values": [] } },
-        "fc2ScalesT": {
-          "dtype": "float32",
-          "shape": [3, 4],
-          "data": { "kind": "linspace", "start": 0.008, "end": 0.037 }
-        }
       },
       "outputs": {
         "outputT": { "dtype": "float32", "shape": [3, 4], "data": { "kind": "constant", "value": 0.0 }, "tolerance": 0 }
@@ -274,7 +261,7 @@
     {
       "name": "q4_colwise_relu_deep_reduction",
       "provenance": {
-        "notes": "A 128-long reduction on both projections, deep enough that splitting it across lanes is worthwhile; the shallow cases above keep the one-thread-per-column schedule."
       },
       "attrs": { "weights_prepacked": 0 },
       "inputs": {
@@ -527,7 +514,7 @@
     {
       "name": "q4_grouped_prefill_partial_tiles",
       "provenance": {
-        "notes": "Half the routed slots the grouped schedule used to demand: 3 experts share 64 slots against 32-row tiles, so most of every tile is padding and the store guard drops more rows than it keeps."
       },
       "attrs": {
         "activation_type": "swiglu",
@@ -573,7 +560,7 @@
     {
       "name": "q8_grouped_prefill_swiglu_blockwise",
       "provenance": {
-        "notes": "Enough routed slots per expert to fill grouped tiles, so one dequantized weight tile is reused across a tile of slots. The tiled reduction sums in a different order than the reference, so the tolerance carries that f32 regrouping; outputs are order 0.4, which a 10% error would move by 0.04."
       },
       "attrs": {
         "activation_type": "swiglu",
@@ -617,6 +604,196 @@
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [48, 64], "tolerance": 0.0003 } }
     }
   ]
 }

 {
   "cases": [
     {
       "name": "q4_colwise_default_relu",
       "provenance": {
+        "notes": "Exercises every exact default at once: 4-bit integer weights, k=1, ReLU, column-wise scales, dense routing, and the raw uint8-packed weight layout."
       },
       "inputs": {
         "inputT": {
     {
       "name": "q4_top1_tie_selects_higher_expert",
       "provenance": {
+        "notes": "Descending `(logit, expert_index)` ordering with selected-weight normalization: equal logits select expert 1, whose output is [16, 32], rather than the all-zero expert 0."
       },
       "attrs": { "normalize_routing_weights": 1 },
       "inputs": {
     {
       "name": "q4_zero_inter",
       "provenance": {
+        "notes": "A zero intermediate dimension has an exact zero result without binding empty weight buffers in a shader. No token, routing or scale value can reach the result once the intermediate width is zero, so the fill values are uniform."
       },
       "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1 },
       "inputs": {
+        "inputT": { "dtype": "float32", "shape": [3, 4], "data": { "kind": "constant", "value": 0.3 } },
+        "routerT": { "dtype": "float32", "shape": [3, 3], "data": { "kind": "constant", "value": 0.7 } },
         "fc1T": { "dtype": "uint8", "shape": [3, 0, 2], "data": { "kind": "values", "values": [] } },
         "fc1ScalesT": { "dtype": "float32", "shape": [3, 0], "data": { "kind": "values", "values": [] } },
         "fc2T": { "dtype": "uint8", "shape": [3, 4, 0], "data": { "kind": "values", "values": [] } },
+        "fc2ScalesT": { "dtype": "float32", "shape": [3, 4], "data": { "kind": "constant", "value": 0.02 } }
       },
       "outputs": {
         "outputT": { "dtype": "float32", "shape": [3, 4], "data": { "kind": "constant", "value": 0.0 }, "tolerance": 0 }
     {
       "name": "q4_colwise_relu_deep_reduction",
       "provenance": {
+        "notes": "A 128-element reduction on both projections exercises the cooperative multi-lane schedule for q4 weights."
       },
       "attrs": { "weights_prepacked": 0 },
       "inputs": {
     {
       "name": "q4_grouped_prefill_partial_tiles",
       "provenance": {
+        "notes": "Three experts share 64 routed slots across 32-row tiles, so most lanes address padding and the store guard discards more rows than it writes."
       },
       "attrs": {
         "activation_type": "swiglu",
     {
       "name": "q8_grouped_prefill_swiglu_blockwise",
       "provenance": {
+        "notes": "Enough routed slots per expert fill grouped tiles, allowing each dequantized weight tile to serve several slots. The tolerance covers f32 reassociation in the tiled reduction while remaining far below a 10% error on outputs near 0.4."
       },
       "attrs": {
         "activation_type": "swiglu",
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [48, 64], "tolerance": 0.0003 } }
+    },
+    {
+      "name": "q4_grouped_stage_t33_h96_i96_block0_swiglu",
+      "attrs": {
+        "activation_type": "swiglu",
+        "expert_weight_bits": 4,
+        "block_size": 0,
+        "k": 2,
+        "normalize_routing_weights": 1,
+        "weights_prepacked": 0,
+        "swiglu_fusion": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [33, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 0.05 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [33, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.44, "cosStep": 0.12, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [3, 192, 48],
+          "data": { "kind": "cycle", "values": [97, 182, 45, 213, 131, 68, 159, 22, 241] }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 96, 48],
+          "data": { "kind": "cycle", "values": [241, 22, 159, 68, 131, 213, 45, 182, 97] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 192],
+          "data": { "kind": "linspace", "start": 0.002, "end": 0.02 }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 96],
+          "data": { "kind": "linspace", "start": 0.003, "end": 0.018 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [33, 96], "tolerance": 0.000002 } },
+      "provenance": {
+        "notes": "Cooperative quantized weight staging across partial 64-column output tiles and routed row tails, with 4/8-bit codes and column/block scales. Small deterministic inputs keep the existing strict f32 tolerance meaningful."
+      }
+    },
+    {
+      "name": "q4_grouped_stage_t31_h96_i32_block16_relu",
+      "attrs": {
+        "activation_type": "relu",
+        "expert_weight_bits": 4,
+        "block_size": 16,
+        "k": 2,
+        "normalize_routing_weights": 1,
+        "weights_prepacked": 0
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [31, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 0.05 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [31, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.44, "cosStep": 0.12, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [3, 32, 48],
+          "data": { "kind": "cycle", "values": [97, 182, 45, 213, 131, 68, 159, 22, 241] }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 96, 16],
+          "data": { "kind": "cycle", "values": [241, 22, 159, 68, 131, 213, 45, 182, 97] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 32, 6],
+          "data": { "kind": "linspace", "start": 0.002, "end": 0.02 }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 96, 2],
+          "data": { "kind": "linspace", "start": 0.003, "end": 0.018 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [31, 96], "tolerance": 0.00005 } },
+      "provenance": {
+        "notes": "Cooperative quantized weight staging across partial 64-column output tiles and routed row tails, with 4/8-bit codes and column/block scales. Small deterministic inputs keep the existing strict f32 tolerance meaningful."
+      }
+    },
+    {
+      "name": "q8_grouped_stage_t65_h128_i96_block32_swiglu",
+      "attrs": {
+        "activation_type": "swiglu",
+        "expert_weight_bits": 8,
+        "block_size": 32,
+        "k": 2,
+        "normalize_routing_weights": 1,
+        "weights_prepacked": 0,
+        "swiglu_fusion": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [65, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 0.05 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [65, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.44, "cosStep": 0.12, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [4, 192, 128],
+          "data": { "kind": "cycle", "values": [97, 182, 45, 213, 131, 68, 159, 22, 241] }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [4, 128, 96],
+          "data": { "kind": "cycle", "values": [241, 22, 159, 68, 131, 213, 45, 182, 97] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [4, 192, 4],
+          "data": { "kind": "linspace", "start": 0.002, "end": 0.02 }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [4, 128, 3],
+          "data": { "kind": "linspace", "start": 0.003, "end": 0.018 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [65, 128], "tolerance": 0.00005 } },
+      "provenance": {
+        "notes": "Cooperative quantized weight staging across partial 64-column output tiles and routed row tails, with 4/8-bit codes and column/block scales. Small deterministic inputs keep the existing strict f32 tolerance meaningful."
+      }
+    },
+    {
+      "name": "q8_grouped_stage_t32_h96_i160_block0_relu",
+      "attrs": {
+        "activation_type": "relu",
+        "expert_weight_bits": 8,
+        "block_size": 0,
+        "k": 1,
+        "normalize_routing_weights": 1,
+        "weights_prepacked": 0
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [32, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 0.05 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [32, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.44, "cosStep": 0.12, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [3, 160, 96],
+          "data": { "kind": "cycle", "values": [97, 182, 45, 213, 131, 68, 159, 22, 241] }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 96, 160],
+          "data": { "kind": "cycle", "values": [241, 22, 159, 68, 131, 213, 45, 182, 97] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 160],
+          "data": { "kind": "linspace", "start": 0.002, "end": 0.02 }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 96],
+          "data": { "kind": "linspace", "start": 0.003, "end": 0.018 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [32, 96], "tolerance": 0.00005 } },
+      "provenance": {
+        "notes": "Cooperative quantized weight staging across partial 64-column output tiles and routed row tails, with 4/8-bit codes and column/block scales. Small deterministic inputs keep the existing strict f32 tolerance meaningful."
+      }
     }
   ]
 }