Xenova HF Staff commited on 7 days ago

Commit

793c2c2

verified ·

1 Parent(s): 07147ca

sync 2e7068faf55e

Browse files

Files changed (17) hide show

README.md +104 -0
build/webgpu/bench.json +693 -0
build/webgpu/expert-group-slots.wgsl.jinja +70 -0
build/webgpu/manifest.json +813 -0
build/webgpu/metadata.json +29 -0
build/webgpu/qmoe-fc1-activation-gemv.wgsl.jinja +167 -0
build/webgpu/qmoe-fc1-activation-grouped-sgmat.wgsl.jinja +228 -0
build/webgpu/qmoe-fc1-activation-grouped.wgsl.jinja +245 -0
build/webgpu/qmoe-fc1-activation-stage.wgsl.jinja +116 -0
build/webgpu/qmoe-mix-stage.wgsl.jinja +33 -0
build/webgpu/qmoe-output-gemv.wgsl.jinja +117 -0
build/webgpu/qmoe-output-grouped-sgmat.wgsl.jinja +196 -0
build/webgpu/qmoe-output-grouped.wgsl.jinja +214 -0
build/webgpu/qmoe-output-stage.wgsl.jinja +78 -0
build/webgpu/qmoe-output-zero-inter.wgsl.jinja +12 -0
build/webgpu/qmoe-route-stage.wgsl.jinja +105 -0
build/webgpu/test.json +622 -0

README.md CHANGED Viewed

@@ -1,3 +1,107 @@
 ---
 license: apache-2.0
 ---

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.QMoE
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 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.
+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
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `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
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32` |
+| `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
+- [`expert-group-slots.wgsl.jinja`](build/webgpu/expert-group-slots.wgsl.jinja)
+- [`qmoe-fc1-activation-gemv.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-gemv.wgsl.jinja)
+- [`qmoe-fc1-activation-grouped-sgmat.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-grouped-sgmat.wgsl.jinja)
+- [`qmoe-fc1-activation-grouped.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-grouped.wgsl.jinja)
+- [`qmoe-fc1-activation-stage.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-stage.wgsl.jinja)
+- [`qmoe-mix-stage.wgsl.jinja`](build/webgpu/qmoe-mix-stage.wgsl.jinja)
+- [`qmoe-output-gemv.wgsl.jinja`](build/webgpu/qmoe-output-gemv.wgsl.jinja)
+- [`qmoe-output-grouped-sgmat.wgsl.jinja`](build/webgpu/qmoe-output-grouped-sgmat.wgsl.jinja)
+- [`qmoe-output-grouped.wgsl.jinja`](build/webgpu/qmoe-output-grouped.wgsl.jinja)
+- [`qmoe-output-stage.wgsl.jinja`](build/webgpu/qmoe-output-stage.wgsl.jinja)
+- [`qmoe-output-zero-inter.wgsl.jinja`](build/webgpu/qmoe-output-zero-inter.wgsl.jinja)
+- [`qmoe-route-stage.wgsl.jinja`](build/webgpu/qmoe-route-stage.wgsl.jinja)
+## 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.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.QMoE", { version: 1 });
+const { outputT } = await kernel({
+  inputT: { data: inputTData, shape: [2, 2, 4] },
+  routerT: { data: routerTData, shape: [4, 2] },
+  fc1T: { data: fc1TData, shape: [2, 4, 2] },
+  fc1ScalesT: { data: fc1ScalesTData, shape: [2, 4] },
+  fc2T: { data: fc2TData, shape: [2, 4, 2] },
+  fc2ScalesT: { data: fc2ScalesTData, shape: [2, 4] },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,693 @@

+{
+  "op": "com.microsoft.QMoE",
+  "tunableSpace": {
+    "workgroupSize": [64, 128, 256],
+    "groupThreads": [8, 16],
+    "groupRegM": [1, 2, 4],
+    "groupRegN": [1, 2, 4],
+    "groupTileK": [8, 16, 32]
+  },
+  "cases": [
+    {
+      "name": "qmoe-q4-relu-t32-h64-e4-i32",
+      "preset": "smoke",
+      "vars": { "tokens": 32, "hidden": 64, "experts": 4, "inter": 32 },
+      "inputs": {
+        "inputT": { "shape": [32, 64], "dtype": "float32", "dist": "normal", "seed": 212, "scale": 0.2 },
+        "routerT": { "shape": [32, 4], "dtype": "float32", "dist": "normal", "seed": 213, "scale": 0.5 },
+        "fc1T": { "shape": [4, 32, 32], "dtype": "uint8", "dist": "uniform", "seed": 214, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [4, 32],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 215,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [4, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 216, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [4, 64],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 217,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [32, 64], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.hidden * args.inter * 2" }]
+      }
+    },
+    {
+      "name": "qmoe-q8-swiglu-t1-h128-e8-i64",
+      "preset": "smoke",
+      "vars": { "tokens": 1, "hidden": 128, "experts": 8, "inter": 64 },
+      "attrs": { "activation_type": "swiglu", "expert_weight_bits": 8, "swiglu_fusion": 1 },
+      "tunables": { "workgroupSize": 64, "decodeLanes": 32, "decodeBlockTarget": 1024 },
+      "tunableSpace": { "decodeMinLaneTrips": [2, 4] },
+      "inputs": {
+        "inputT": { "shape": [1, 128], "dtype": "float32", "dist": "normal", "seed": 312, "scale": 0.2 },
+        "routerT": { "shape": [1, 8], "dtype": "float32", "dist": "normal", "seed": 313, "scale": 0.5 },
+        "fc1T": { "shape": [8, 128, 128], "dtype": "uint8", "dist": "uniform", "seed": 314, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [8, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 315,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [8, 128, 64], "dtype": "uint8", "dist": "uniform", "seed": 316, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [8, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 317,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [1, 128], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "qmoe-q4-swiglu-t1-h2048-e8-k2-i768-decode",
+      "preset": "smoke",
+      "vars": { "tokens": 1, "hidden": 2048, "experts": 8, "inter": 768, "topK": 2, "bits": 4 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 2 },
+      "tunableSpace": { "workgroupSize": [64, 128, 256], "decodeLanes": [8, 16, 32, 64] },
+      "inputs": {
+        "inputT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 400, "scale": 0.2 },
+        "routerT": { "shape": [1, 8], "dtype": "float32", "dist": "normal", "seed": 401, "scale": 0.5 },
+        "fc1T": { "shape": [8, 1536, 1024], "dtype": "uint8", "dist": "uniform", "seed": 402, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [8, 1536],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 403,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [8, 2048, 384], "dtype": "uint8", "dist": "uniform", "seed": 404, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [8, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 405,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [1, 2048], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "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."
+      }
+    },
+    {
+      "name": "qmoe-q4-swiglu-t1-h1024-e16-k4-i512-decode-topk4",
+      "preset": "stress",
+      "vars": { "tokens": 1, "hidden": 1024, "experts": 16, "inter": 512, "topK": 4, "bits": 4 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 4 },
+      "inputs": {
+        "inputT": { "shape": [1, 1024], "dtype": "float32", "dist": "normal", "seed": 420, "scale": 0.2 },
+        "routerT": { "shape": [1, 16], "dtype": "float32", "dist": "normal", "seed": 421, "scale": 0.5 },
+        "fc1T": { "shape": [16, 1024, 512], "dtype": "uint8", "dist": "uniform", "seed": 422, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [16, 1024],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 423,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [16, 1024, 256], "dtype": "uint8", "dist": "uniform", "seed": 424, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [16, 1024],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 425,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [1, 1024], "dtype": "float32" } },
+      "bench": {
+        "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."
+      }
+    },
+    {
+      "name": "qmoe-q4-swiglu-t8-h2048-e8-k2-i768-decode-batch8",
+      "preset": "stress",
+      "vars": { "tokens": 8, "hidden": 2048, "experts": 8, "inter": 768, "topK": 2, "bits": 4 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 2 },
+      "tunableSpace": { "decodeBlockTarget": [128, 256] },
+      "inputs": {
+        "inputT": { "shape": [8, 2048], "dtype": "float32", "dist": "normal", "seed": 440, "scale": 0.2 },
+        "routerT": { "shape": [8, 8], "dtype": "float32", "dist": "normal", "seed": 441, "scale": 0.5 },
+        "fc1T": { "shape": [8, 1536, 1024], "dtype": "uint8", "dist": "uniform", "seed": 442, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [8, 1536],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 443,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [8, 2048, 384], "dtype": "uint8", "dist": "uniform", "seed": 444, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [8, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 445,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [8, 2048], "dtype": "float32" } },
+      "bench": {
+        "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."
+      }
+    },
+    {
+      "name": "qmoe-q8-swiglu-t1-h1024-e8-k2-i1024-decode-q8",
+      "preset": "stress",
+      "vars": { "tokens": 1, "hidden": 1024, "experts": 8, "inter": 1024, "topK": 2, "bits": 8 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "expert_weight_bits": 8, "k": 2 },
+      "inputs": {
+        "inputT": { "shape": [1, 1024], "dtype": "float32", "dist": "normal", "seed": 460, "scale": 0.2 },
+        "routerT": { "shape": [1, 8], "dtype": "float32", "dist": "normal", "seed": 461, "scale": 0.5 },
+        "fc1T": { "shape": [8, 2048, 1024], "dtype": "uint8", "dist": "uniform", "seed": 462, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [8, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 463,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [8, 1024, 1024], "dtype": "uint8", "dist": "uniform", "seed": 464, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [8, 1024],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 465,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [1, 1024], "dtype": "float32" } },
+      "bench": {
+        "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."
+      }
+    },
+    {
+      "name": "qmoe-q4-relu-t1-h2048-e8-i2048-decode-relu",
+      "preset": "stress",
+      "vars": { "tokens": 1, "hidden": 2048, "experts": 8, "inter": 2048, "topK": 1, "bits": 4 },
+      "inputs": {
+        "inputT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 480, "scale": 0.2 },
+        "routerT": { "shape": [1, 8], "dtype": "float32", "dist": "normal", "seed": 481, "scale": 0.5 },
+        "fc1T": { "shape": [8, 2048, 1024], "dtype": "uint8", "dist": "uniform", "seed": 482, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [8, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 483,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [8, 2048, 1024], "dtype": "uint8", "dist": "uniform", "seed": 484, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [8, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 485,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [1, 2048], "dtype": "float32" } },
+      "bench": {
+        "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."
+      }
+    },
+    {
+      "name": "qmoe-q4-swiglu-t48-h1024-e8-k2-i1024-prefill-band",
+      "preset": "stress",
+      "vars": { "tokens": 48, "hidden": 1024, "experts": 8, "inter": 1024, "topK": 2 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 2 },
+      "inputs": {
+        "inputT": { "shape": [48, 1024], "dtype": "float32", "dist": "normal", "seed": 600, "scale": 0.2 },
+        "routerT": { "shape": [48, 8], "dtype": "float32", "dist": "normal", "seed": 601, "scale": 0.5 },
+        "fc1T": { "shape": [8, 2048, 512], "dtype": "uint8", "dist": "uniform", "seed": 502, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [8, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 503,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [8, 1024, 512], "dtype": "uint8", "dist": "uniform", "seed": 504, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [8, 1024],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 505,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [48, 1024], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "qmoe-q4-swiglu-t96-h1024-e8-k2-i1024-prefill-small",
+      "preset": "stress",
+      "vars": { "tokens": 96, "hidden": 1024, "experts": 8, "inter": 1024, "topK": 2 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 2 },
+      "tunables": { "workgroupSize": 128, "decodeLanes": 32, "decodeMinLaneTrips": 4 },
+      "inputs": {
+        "inputT": { "shape": [96, 1024], "dtype": "float32", "dist": "normal", "seed": 620, "scale": 0.2 },
+        "routerT": { "shape": [96, 8], "dtype": "float32", "dist": "normal", "seed": 621, "scale": 0.5 },
+        "fc1T": { "shape": [8, 2048, 512], "dtype": "uint8", "dist": "uniform", "seed": 502, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [8, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 503,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [8, 1024, 512], "dtype": "uint8", "dist": "uniform", "seed": 504, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [8, 1024],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 505,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [96, 1024], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "qmoe-q4-swiglu-t128-h1024-e8-k2-i1024-prefill",
+      "preset": "stress",
+      "vars": { "tokens": 128, "hidden": 1024, "experts": 8, "inter": 1024, "topK": 2 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 2 },
+      "tunables": { "workgroupSize": 128, "decodeLanes": 32, "decodeMinLaneTrips": 4 },
+      "tunableSpace": { "groupRegM": [2, 4], "groupRegN": [2, 4], "groupTileK": [8, 16] },
+      "inputs": {
+        "inputT": { "shape": [128, 1024], "dtype": "float32", "dist": "normal", "seed": 500, "scale": 0.2 },
+        "routerT": { "shape": [128, 8], "dtype": "float32", "dist": "normal", "seed": 501, "scale": 0.5 },
+        "fc1T": { "shape": [8, 2048, 512], "dtype": "uint8", "dist": "uniform", "seed": 502, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [8, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 503,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [8, 1024, 512], "dtype": "uint8", "dist": "uniform", "seed": 504, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [8, 1024],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 505,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [128, 1024], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "qmoe-q4-swiglu-t512-h1024-e8-k2-i1024-prefill-long",
+      "preset": "stress",
+      "vars": { "tokens": 512, "hidden": 1024, "experts": 8, "inter": 1024, "topK": 2 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 2 },
+      "inputs": {
+        "inputT": { "shape": [512, 1024], "dtype": "float32", "dist": "normal", "seed": 520, "scale": 0.2 },
+        "routerT": { "shape": [512, 8], "dtype": "float32", "dist": "normal", "seed": 521, "scale": 0.5 },
+        "fc1T": { "shape": [8, 2048, 512], "dtype": "uint8", "dist": "uniform", "seed": 522, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [8, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 523,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [8, 1024, 512], "dtype": "uint8", "dist": "uniform", "seed": 524, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [8, 1024],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 525,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [512, 1024], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "qmoe-q4-swiglu-t384-h1024-e32-k1-i512-prefill-band",
+      "preset": "stress",
+      "vars": { "tokens": 384, "hidden": 1024, "experts": 32, "inter": 512, "topK": 1 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 1 },
+      "inputs": {
+        "inputT": { "shape": [384, 1024], "dtype": "float32", "dist": "normal", "seed": 700, "scale": 0.2 },
+        "routerT": { "shape": [384, 32], "dtype": "float32", "dist": "normal", "seed": 701, "scale": 0.5 },
+        "fc1T": { "shape": [32, 1024, 512], "dtype": "uint8", "dist": "uniform", "seed": 702, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [32, 1024],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 703,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [32, 1024, 256], "dtype": "uint8", "dist": "uniform", "seed": 704, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [32, 1024],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 705,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [384, 1024], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "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 },
+      "inputs": {
+        "inputT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 4011, "scale": 0.2 },
+        "routerT": { "shape": [1, 64], "dtype": "float32", "dist": "normal", "seed": 4012, "scale": 0.5 },
+        "fc1T": { "shape": [64, 1536, 1024], "dtype": "uint8", "dist": "uniform", "seed": 4013, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [64, 1536],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4014,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [64, 2048, 384], "dtype": "uint8", "dist": "uniform", "seed": 4015, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [64, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4016,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [1, 2048], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
+      }
+    },
+    {
+      "name": "qmoe-q4-qwen3-moe-prefill-t256-h2048-e64-k8-i768",
+      "preset": "model",
+      "provenance": {
+        "notes": "Qwen3-MoE class defaults at a 256-token prefill chunk; 2048 routed slots clear the grouped schedule's slot floor for 64 experts."
+      },
+      "vars": { "tokens": 256, "hidden": 2048, "experts": 64, "inter": 768, "topK": 8 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 8 },
+      "inputs": {
+        "inputT": { "shape": [256, 2048], "dtype": "float32", "dist": "normal", "seed": 4021, "scale": 0.2 },
+        "routerT": { "shape": [256, 64], "dtype": "float32", "dist": "normal", "seed": 4022, "scale": 0.5 },
+        "fc1T": { "shape": [64, 1536, 1024], "dtype": "uint8", "dist": "uniform", "seed": 4023, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [64, 1536],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4024,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [64, 2048, 384], "dtype": "uint8", "dist": "uniform", "seed": 4025, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [64, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4026,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [256, 2048], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "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 },
+      "inputs": {
+        "inputT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 4031, "scale": 0.2 },
+        "routerT": { "shape": [1, 32], "dtype": "float32", "dist": "normal", "seed": 4032, "scale": 0.5 },
+        "fc1T": { "shape": [32, 4096, 1024], "dtype": "uint8", "dist": "uniform", "seed": 4033, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [32, 4096],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4034,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [32, 2048, 1024], "dtype": "uint8", "dist": "uniform", "seed": 4035, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [32, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4036,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [1, 2048], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
+      }
+    },
+    {
+      "name": "qmoe-q4-olmoe-prefill-t128-h2048-e32-k8-i2048",
+      "preset": "model",
+      "provenance": { "notes": "OLMoE class defaults at a 128-token prefill chunk." },
+      "vars": { "tokens": 128, "hidden": 2048, "experts": 32, "inter": 2048, "topK": 8 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 8 },
+      "inputs": {
+        "inputT": { "shape": [128, 2048], "dtype": "float32", "dist": "normal", "seed": 4041, "scale": 0.2 },
+        "routerT": { "shape": [128, 32], "dtype": "float32", "dist": "normal", "seed": 4042, "scale": 0.5 },
+        "fc1T": { "shape": [32, 4096, 1024], "dtype": "uint8", "dist": "uniform", "seed": 4043, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [32, 4096],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4044,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [32, 2048, 1024], "dtype": "uint8", "dist": "uniform", "seed": 4045, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [32, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4046,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [128, 2048], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "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 },
+      "inputs": {
+        "inputT": { "shape": [1, 2880], "dtype": "float32", "dist": "normal", "seed": 4051, "scale": 0.2 },
+        "routerT": { "shape": [1, 16], "dtype": "float32", "dist": "normal", "seed": 4052, "scale": 0.5 },
+        "fc1T": { "shape": [16, 5760, 1440], "dtype": "uint8", "dist": "uniform", "seed": 4053, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [16, 5760],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4054,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [16, 2880, 1440], "dtype": "uint8", "dist": "uniform", "seed": 4055, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [16, 2880],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4056,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [1, 2880], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
+      }
+    },
+    {
+      "name": "qmoe-q4-gpt-oss-prefill-t256-h2880-e16-k4-i2880",
+      "preset": "model",
+      "provenance": { "notes": "GPT-OSS class defaults at a 256-token prefill chunk." },
+      "vars": { "tokens": 256, "hidden": 2880, "experts": 16, "inter": 2880, "topK": 4 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 4 },
+      "inputs": {
+        "inputT": { "shape": [256, 2880], "dtype": "float32", "dist": "normal", "seed": 4061, "scale": 0.2 },
+        "routerT": { "shape": [256, 16], "dtype": "float32", "dist": "normal", "seed": 4062, "scale": 0.5 },
+        "fc1T": { "shape": [16, 5760, 1440], "dtype": "uint8", "dist": "uniform", "seed": 4063, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [16, 5760],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4064,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [16, 2880, 1440], "dtype": "uint8", "dist": "uniform", "seed": 4065, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [16, 2880],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4066,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [256, 2880], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "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 },
+      "inputs": {
+        "inputT": { "shape": [1, 4096], "dtype": "float32", "dist": "normal", "seed": 4071, "scale": 0.2 },
+        "routerT": { "shape": [1, 4], "dtype": "float32", "dist": "normal", "seed": 4072, "scale": 0.5 },
+        "fc1T": { "shape": [4, 12800, 2048], "dtype": "uint8", "dist": "uniform", "seed": 4073, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [4, 12800],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4074,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [4, 4096, 3200], "dtype": "uint8", "dist": "uniform", "seed": 4075, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [4, 4096],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4076,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [1, 4096], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * args.bits / 2" }]
+      }
+    },
+    {
+      "name": "qmoe-q4-phimoe-prefill-t256-h4096-e4-k2-i6400",
+      "preset": "model",
+      "provenance": { "notes": "Phi-MoE class defaults at a 256-token prefill chunk." },
+      "vars": { "tokens": 256, "hidden": 4096, "experts": 4, "inter": 6400, "topK": 2 },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 1, "k": 2 },
+      "inputs": {
+        "inputT": { "shape": [256, 4096], "dtype": "float32", "dist": "normal", "seed": 4081, "scale": 0.2 },
+        "routerT": { "shape": [256, 4], "dtype": "float32", "dist": "normal", "seed": 4082, "scale": 0.5 },
+        "fc1T": { "shape": [4, 12800, 2048], "dtype": "uint8", "dist": "uniform", "seed": 4083, "min": 0, "max": 255 },
+        "fc1ScalesT": {
+          "shape": [4, 12800],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4084,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "fc2T": { "shape": [4, 4096, 3200], "dtype": "uint8", "dist": "uniform", "seed": 4085, "min": 0, "max": 255 },
+        "fc2ScalesT": {
+          "shape": [4, 4096],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 4086,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": { "outputT": { "shape": [256, 4096], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    }
+  ]
+}

build/webgpu/expert-group-slots.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,70 @@

+{{ env.wgsl.resourceDeclarations }}
+// Group this token chunk's routed slots by expert, then cut each expert's run into
+// MTILE-row tiles for the grouped FFN stage.
+//
+// The unit being grouped is the SLOT, not the token: a token contributes TOP_K slots and each
+// one carries its own expert, so the slot space is a K-COVER of [0, chunkTokens * TOP_K) rather
+// than a partition of the tokens. That is what lets the grouped tiles serve any k -- a token
+// simply appears once in each of the (up to TOP_K) expert runs it was routed to.
+//
+// tile_meta[0] is the tile count; tile t occupies the three words at 1 + t * 3 and holds its
+// expert, the start of its slice of slot_list, and how many rows of that slice are real.
+//
+// Each expert run may end with one partial tile, so the grouped schedule carries
+// about EXPERTS * MTILE / 2 padded rows.
+const EXPERTS: u32 = {{ experts }}u;
+const TOP_K: u32 = {{ topK }}u;
+const MTILE: u32 = {{ groupTileM }}u;
+const WG: u32 = {{ groupRouteWorkgroup }}u;
+var<workgroup> counts: array<atomic<u32>, {{ experts }}>;
+var<workgroup> offsets: array<u32, {{ experts }}>;
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(local_invocation_index) lid: u32) {
+  let slot_total = params.tokenCount * TOP_K;
+  let slot_base = params.tokenOffset * TOP_K;
+  for (var e = lid; e < EXPERTS; e = e + WG) {
+    atomicStore(&counts[e], 0u);
+  }
+  workgroupBarrier();
+  for (var p = lid; p < slot_total; p = p + WG) {
+    atomicAdd(&counts[route_expert[slot_base + p]], 1u);
+  }
+  workgroupBarrier();
+  // One lane lays out the EXPERTS-sized run and tile tables. The chunk-sized
+  // scatter below is distributed over the workgroup.
+  if (lid == 0u) {
+    var offset = 0u;
+    var tile = 0u;
+    for (var e = 0u; e < EXPERTS; e = e + 1u) {
+      offsets[e] = offset;
+      let n = atomicLoad(&counts[e]);
+      var done = 0u;
+      while (done < n) {
+        let rows = min(MTILE, n - done);
+        tile_meta[1u + tile * 3u] = e;
+        tile_meta[1u + tile * 3u + 1u] = offset + done;
+        tile_meta[1u + tile * 3u + 2u] = rows;
+        tile = tile + 1u;
+        done = done + rows;
+      }
+      offset = offset + n;
+      atomicStore(&counts[e], 0u);
+    }
+    tile_meta[0] = tile;
+  }
+  workgroupBarrier();
+  // Each slot atomically claims a position in its expert's run. Arrival order is
+  // safe because downstream dots use the slot's own row and reduce to its token;
+  // no consumer assigns meaning to order within a run.
+  for (var p = lid; p < slot_total; p = p + WG) {
+    let expert = route_expert[slot_base + p];
+    slot_list[offsets[expert] + atomicAdd(&counts[expert], 1u)] = p;
+  }
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,813 @@

+{
+  "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"] },
+    "expert_weight_bits": { "values": [4, 8] },
+    "normalize_routing_weights": { "values": [0, 1] },
+    "quant_type": { "values": ["int"] },
+    "swiglu_fusion": { "values": [0, 1] },
+    "use_sparse_mixer": { "values": [0] },
+    "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)",
+    "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))))",
+    "hiddenChunkCount": "max(1, ceilDiv(numTokens, max(1, hiddenChunkTokens)))",
+    "routeScratchBytes": "numTokens * topK * 4",
+    "routedScratchFits": "routeScratchBytes <= device.limits.maxStorageBufferBindingSize and routeScratchBytes <= device.limits.maxBufferSize",
+    "groupTileM": "tunables.groupThreads * tunables.groupRegM",
+    "groupTileN": "tunables.groupThreads * tunables.groupRegN",
+    "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": [
+    {
+      "id": "quant_zero_inter",
+      "priority": 20,
+      "when": ["quantContract", "workgroupSizeOk", "interSize == 0"],
+      "passes": [
+        {
+          "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" }
+        }
+      ]
+    },
+    {
+      "id": "quant_grouped_sgmat_routed",
+      "priority": 32,
+      "when": ["groupedSgmatOk"],
+      "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]" },
+        { "id": "hiddenAct", "dtype": "float32", "shape": "[hiddenChunkTokens * topK * interSize]" },
+        { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" },
+        { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" },
+        { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenSize)]" }
+      ],
+      "passes": [
+        {
+          "id": "route_stage",
+          "name": "QMoE.RouteStage",
+          "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": "route",
+          "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" }
+        },
+        {
+          "id": "token_chunks",
+          "repeat": { "count": "hiddenChunkCount", "index": "chunk" },
+          "passes": [
+            {
+              "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 },
+        "subgroupMinSize": 32,
+        "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",
+        "groupRouteWorkgroup": "tunables.groupRouteWorkgroup"
+      },
+      "intermediates": [
+        { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" },
+        { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" },
+        { "id": "hiddenAct", "dtype": "float32", "shape": "[hiddenChunkTokens * topK * interSize]" },
+        { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" },
+        { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" },
+        { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenSize)]" }
+      ],
+      "passes": [
+        {
+          "id": "route_stage",
+          "name": "QMoE.RouteStage",
+          "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": "route",
+          "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" }
+        },
+        {
+          "id": "token_chunks",
+          "repeat": { "count": "hiddenChunkCount", "index": "chunk" },
+          "passes": [
+            {
+              "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"
+              }
+            }
+          ]
+        }
+      ]
+    },
+    {
+      "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]" },
+        { "id": "hiddenAct", "dtype": "float32", "shape": "[hiddenChunkTokens * topK * interSize]" }
+      ],
+      "passes": [
+        {
+          "id": "route_stage",
+          "name": "QMoE.RouteStage",
+          "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": "route",
+          "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" }
+        },
+        {
+          "id": "token_chunks",
+          "repeat": { "count": "hiddenChunkCount", "index": "chunk" },
+          "passes": [
+            {
+              "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"
+              }
+            },
+            {
+              "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)"
+              }
+            }
+          ]
+        }
+      ]
+    },
+    {
+      "id": "quant_split_routed",
+      "priority": 10,
+      "when": ["quantContract", "workgroupSizeOk", "hiddenChunkFits", "routedScratchFits", "interSize > 0"],
+      "intermediates": [
+        { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" },
+        { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" },
+        { "id": "hiddenAct", "dtype": "float32", "shape": "[hiddenChunkTokens * topK * interSize]" }
+      ],
+      "passes": [
+        {
+          "id": "route_stage",
+          "name": "QMoE.RouteStage",
+          "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": "route",
+          "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" }
+        },
+        {
+          "id": "token_chunks",
+          "repeat": { "count": "hiddenChunkCount", "index": "chunk" },
+          "passes": [
+            {
+              "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"
+              }
+            }
+          ]
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,29 @@

+{
+  "name": "com.microsoft.QMoE",
+  "id": "_com_microsoft_qmoe_webgpu_9a12372",
+  "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": "hZMDKX6e3562/hNqPc5QSpMZ7BRW0wcSrH63e8VTdM4=",
+      "qmoe-fc1-activation-gemv.wgsl.jinja": "UvP8Vyx1E9+/+I09w/z9iB6h6MHrSrOmjKfIelyIS0M=",
+      "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja": "bRKPszdbDTG0DlHeiiQ1lxyjJUJz0JVt9WZDsv8KVg8=",
+      "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": "XaFjSitK6tkK/QkrQ6wRConJLNekx+4u/FpxuFmh7DU=",
+      "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": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.QMoE" }
+}

build/webgpu/qmoe-fc1-activation-gemv.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,167 @@

+{{ env.wgsl.resourceDeclarations }}
+const HIDDEN: u32 = {{ hidden }}u;
+const FC1_ROWS: u32 = {{ fc1Rows }}u;
+const INTER: u32 = {{ inter }}u;
+const TOP_K: u32 = {{ topK }}u;
+// The split stage hands one thread a whole HIDDEN-long reduction, so a decode
+// dispatch offers only tokens * TOP_K * INTER invocations. Here a workgroup owns
+// ROWS intermediate columns of one routed slot and LANES lanes split the
+// reduction, which is LANES times the invocations for the same arithmetic.
+const LANES: u32 = {{ decodeLanes }}u;
+const ROWS: u32 = {{ decodeRows }}u;
+// Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
+// weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
+// first along the reduction axis, two codes per stored word.
+const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const FC1_PACKED_COLS: u32 = {{ fc1PackedCols }}u;
+{% if not colWiseScales %}
+// Block-wise scales (fc1_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
+// axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
+const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const FC1_BLOCKS: u32 = {{ fc1ScaleBlocks }}u;
+{% endif %}
+fn fc1_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * FC1_ROWS + row;
+{% if packSize == 2 %}
+  let word = fc1_experts_weights[base * FC1_PACKED_COLS + col / 2u];
+  let code = f32((word >> ((col & 1u) * 4u)) & 15u);
+{% else %}
+  let code = f32(fc1_experts_weights[base * FC1_PACKED_COLS + col] & 255u);
+{% endif %}
+{% if not colWiseScales %}
+  return (code - MIDPOINT) * fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK];
+{% else %}
+  return (code - MIDPOINT) * fc1_scales[base];
+{% endif %}
+}
+// Decode two consecutive reduction-axis codes together. Four-bit codes share a
+// stored word; eight-bit codes occupy adjacent words.
+fn fc1_weight_pair(expert: u32, row: u32, col: u32) -> vec2<f32> {
+  let base = expert * FC1_ROWS + row;
+{% if packSize == 2 %}
+  let word = fc1_experts_weights[base * FC1_PACKED_COLS + col / 2u];
+  let lo = f32(word & 15u) - MIDPOINT;
+  let hi = f32((word >> 4u) & 15u) - MIDPOINT;
+{% else %}
+  let lo = f32(fc1_experts_weights[base * FC1_PACKED_COLS + col] & 255u) - MIDPOINT;
+  let hi = f32(fc1_experts_weights[base * FC1_PACKED_COLS + col + 1u] & 255u) - MIDPOINT;
+{% endif %}
+{% if not colWiseScales %}
+  return vec2<f32>(lo * fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK],
+                   hi * fc1_scales[base * FC1_BLOCKS + (col + 1u) / QUANT_BLOCK]);
+{% else %}
+  let scale = fc1_scales[base];
+  return vec2<f32>(lo * scale, hi * scale);
+{% endif %}
+}
+{% if activationType == "swiglu" %}
+fn activate(gate_in: f32, up_in: f32) -> f32 {
+{% if hasSwigluLimit %}
+  // swiglu_limit clamps gate from above and up to [-limit, limit] before the
+  // product; when absent, no clamp is applied.
+  let gate = min(gate_in, {{ swigluLimit }});
+  let up = clamp(up_in, -({{ swigluLimit }}), {{ swigluLimit }});
+{% else %}
+  let gate = gate_in;
+  let up = up_in;
+{% endif %}
+  return gate / (1.0 + exp(-params.activationAlpha * gate)) * (up + params.activationBeta);
+}
+{% else %}
+fn activate(value: f32) -> f32 {
+  return max(value, 0.0);
+}
+{% endif %}
+{% macro rowlane_fold(arrays, lanes="LANES", lane="lane", slot="slot") %}
+  var n = {{ lanes }};
+  while (n > 1u) {
+    let half = (n + 1u) / 2u;
+    if ({{ lane }} + half < n) {
+{% for a in arrays %}
+      {{ a }}[{{ slot }}] = {{ a }}[{{ slot }}] + {{ a }}[{{ slot }} + half];
+{% endfor %}
+    }
+    workgroupBarrier();
+    n = half;
+  }
+{%- endmacro %}
+var<workgroup> gate_partial: array<f32, {{ decodeLanes * decodeRows }}>;
+{% if activationType == "swiglu" %}
+// The gate and up rows of an interleaved SwiGLU pair share the same activation
+// value, so one workgroup reduces both and folds them together.
+var<workgroup> up_partial: array<f32, {{ decodeLanes * decodeRows }}>;
+{% endif %}
+@compute @workgroup_size(LANES, ROWS, 1)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  let i = wid.x * ROWS + lid.y;
+  let lane = lid.x;
+  // wid.y is the routed slot's index within this token chunk -- the split
+  // stage's `local_token * TOP_K + j` -- so hidden_act is addressed identically
+  // by both schedules.
+  let local_slot = wid.y;
+  let token = params.tokenOffset + local_slot / TOP_K;
+  let in_range = i < INTER;
+  var gate = 0.0;
+{% if activationType == "swiglu" %}
+  var up = 0.0;
+{% endif %}
+  if (in_range) {
+    let expert = route_expert[token * TOP_K + local_slot % TOP_K];
+    // Two reduction columns per step: at 4 bits they are one stored word.
+    for (var c = lane * 2u; c + 1u < HIDDEN; c = c + LANES * 2u) {
+      let x0 = f32(input[token * HIDDEN + c]);
+      let x1 = f32(input[token * HIDDEN + c + 1u]);
+{% if activationType == "swiglu" %}
+      let g = fc1_weight_pair(expert, i * 2u, c);
+      let u = fc1_weight_pair(expert, i * 2u + 1u, c);
+      gate = gate + x0 * g.x + x1 * g.y;
+      up = up + x0 * u.x + x1 * u.y;
+{% else %}
+      let g = fc1_weight_pair(expert, i, c);
+      gate = gate + x0 * g.x + x1 * g.y;
+{% endif %}
+    }
+    // An odd reduction length leaves one column, which only 8-bit weights can have.
+    if ((HIDDEN & 1u) == 1u && lane == 0u) {
+      let h = HIDDEN - 1u;
+      let x = f32(input[token * HIDDEN + h]);
+{% if activationType == "swiglu" %}
+      gate = gate + x * fc1_weight(expert, i * 2u, h);
+      up = up + x * fc1_weight(expert, i * 2u + 1u, h);
+{% else %}
+      gate = gate + x * fc1_weight(expert, i, h);
+{% endif %}
+    }
+  }
+  let slot = lid.y * LANES + lane;
+  gate_partial[slot] = gate;
+{% if activationType == "swiglu" %}
+  up_partial[slot] = up;
+{% endif %}
+  workgroupBarrier();
+{% if activationType == "swiglu" %}
+{{ rowlane_fold(["gate_partial", "up_partial"]) }}
+{% else %}
+{{ rowlane_fold(["gate_partial"]) }}
+{% endif %}
+  if (lane == 0u && in_range) {
+    let base = lid.y * LANES;
+{% if activationType == "swiglu" %}
+    hidden_act[local_slot * INTER + i] = activate(gate_partial[base], up_partial[base]);
+{% else %}
+    hidden_act[local_slot * INTER + i] = activate(gate_partial[base]);
+{% endif %}
+  }
+}

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

	@@ -0,0 +1,228 @@

+enable subgroups;
+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
+// lane holds are already an even/odd pair, so the epilogue needs no reshuffle.
+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;
+const SUB_ROWS: u32 = 16u;
+{% set isSwiglu = activationType == "swiglu" %}
+// Weight rows per output column: the fused gate/up pair, or one.
+const FUSION: u32 = {{ 2 if isSwiglu else 1 }}u;
+// Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
+// weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
+// first along the reduction axis, two codes per stored word.
+const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const FC1_PACKED_COLS: u32 = {{ fc1PackedCols }}u;
+{% if not colWiseScales %}
+// Block-wise scales (fc1_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
+// axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
+const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const FC1_BLOCKS: u32 = {{ fc1ScaleBlocks }}u;
+{% endif %}
+// 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 %}
+  let pair = base * FC1_PACKED_COLS + col / 2u;
+  let w0 = fc1_experts_weights[pair];
+  let w1 = fc1_experts_weights[pair + 1u];
+  let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
+                        f32(w1 & 15u), f32((w1 >> 4u) & 15u));
+{% else %}
+  let word = base * FC1_PACKED_COLS + col;
+  let codes = vec4<f32>(f32(fc1_experts_weights[word] & 255u), f32(fc1_experts_weights[word + 1u] & 255u),
+                        f32(fc1_experts_weights[word + 2u] & 255u), f32(fc1_experts_weights[word + 3u] & 255u));
+{% endif %}
+{% if not colWiseScales %}
+  let scale = fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK];
+{% else %}
+  let scale = fc1_scales[base];
+{% endif %}
+  return (codes - vec4<f32>(MIDPOINT)) * scale;
+}
+{% if activationType == "swiglu" %}
+fn activate(gate_in: f32, up_in: f32) -> f32 {
+{% if hasSwigluLimit %}
+  // swiglu_limit clamps gate from above and up to [-limit, limit] before the
+  // product; when absent, no clamp is applied.
+  let gate = min(gate_in, {{ swigluLimit }});
+  let up = clamp(up_in, -({{ swigluLimit }}), {{ swigluLimit }});
+{% else %}
+  let gate = gate_in;
+  let up = up_in;
+{% endif %}
+  return gate / (1.0 + exp(-params.activationAlpha * gate)) * (up + params.activationBeta);
+}
+{% else %}
+fn activate(value: f32) -> f32 {
+  return max(value, 0.0);
+}
+{% 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]);
+    }
+  } else {
+    for (var i = 0u; i < 8u; i = i + 1u) {
+      tile_A[row * TILE_K + col + i] = 0.0;
+    }
+  }
+}
+// `row` indexes weight rows, so under SwiGLU it walks gate and up alternately. Rows whose
+// output column is past INTER stage zero: they contribute nothing, are never stored, and
+// their packed reads would be out of range.
+fn loadSHMB(expert: u32, weight_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+  let n = weight_base + row;
+  let col = c_idx * 16u;
+  let live = n / FUSION < INTER;
+  for (var i = 0u; i < 16u; i = i + 4u) {
+    var quad = vec4<f32>(0.0);
+    if (live) {
+      quad = fc1_weight_vec4(expert, n, k_idx + col + i);
+    }
+    let slot = row * TILE_K + col + i;
+    tile_B[slot] = quad.x;
+    tile_B[slot + 1u] = quad.y;
+    tile_B[slot + 2u] = quad.z;
+    tile_B[slot + 3u] = quad.w;
+  }
+}
+@compute @workgroup_size({{ sgmatWorkgroup }}, 1, 1)
+fn main(
+  @builtin(workgroup_id) wid: vec3<u32>,
+  @builtin(local_invocation_index) local_idx: u32,
+  @builtin(subgroup_invocation_id) sg_id: u32,
+  @builtin(subgroup_size) sg_size: u32
+) {
+  let tile = wid.x;
+  if (tile >= tile_meta[0]) {
+    return;
+  }
+  let expert = tile_meta[1u + tile * 3u];
+  let slice_base = tile_meta[1u + tile * 3u + 1u];
+  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;
+  let subtile_idy = subtile_id % {{ sgmatRowSubtiles }}u;
+  let base_A = subtile_idy * SUB_ROWS;
+  let base_B = subtile_idx * SUB_COLS;
+{% for row in [0, 1] %}{% for col in [0, 1, 2, 3] %}
+  var matC{{ row }}{{ col }}: subgroup_matrix_result<f32, 8, 8>;
+{% 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>>(&tile_A, matrix_a_offset, false, TILE_K);
+      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, 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>>(&tile_B, matrix_b_offset, true, TILE_K);
+      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
+      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
+      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
+      matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
+      matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
+      matC02 = subgroupMatrixMultiplyAccumulate(matA0, matB2, matC02);
+      matC03 = subgroupMatrixMultiplyAccumulate(matA0, matB3, matC03);
+      matC10 = subgroupMatrixMultiplyAccumulate(matA1, matB0, matC10);
+      matC11 = subgroupMatrixMultiplyAccumulate(matA1, matB1, matC11);
+      matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12);
+      matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13);
+    }
+    workgroupBarrier();
+  }
+  // A lane holds result columns `col` and `col2 = col + 1` of each 8-wide block, and `col`
+  // is even -- so under SwiGLU that pair is exactly one output column's gate and up.
+  let row = sg_id / 4u;
+  let col = (sg_id % 4u) * 2u;
+  let col2 = col + 1u;
+  let out_weight_base = weight_base + base_B;
+{% for group in [0, 1] %}
+  // 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(&scratch[subtile_id][{{ c }}], 0u, matC{{ group }}{{ c }}, false, 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 %}
+{% for half in ['col', 'col2'] %}
+      {
+        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 %}
+{% endif %}
+{% endfor %}
+    }
+  }
+  workgroupBarrier();
+{% endfor %}
+}

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

	@@ -0,0 +1,245 @@

+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.QMoE, expert-grouped FC1 + activation stage.
+// The group stage has sorted routed slots by expert, so an MTILE-row tile shares
+// one expert. Its weight tile is dequantized once into workgroup memory and
+// reused across every row.
+const HIDDEN: u32 = {{ hidden }}u;
+const FC1_ROWS: u32 = {{ fc1Rows }}u;
+const INTER: u32 = {{ inter }}u;
+const TOP_K: u32 = {{ topK }}u;
+const MTILE: u32 = {{ groupTileM }}u;
+const NTILE: u32 = {{ groupTileN }}u;
+const KTILE: u32 = {{ groupTileK }}u;
+const THREADS_SIDE: u32 = {{ groupThreads }}u;
+const THREADS: u32 = THREADS_SIDE * THREADS_SIDE;
+{% set isSwiglu = activationType == "swiglu" %}
+{% set primaryRow = "col * 2u" if isSwiglu else "col" %}
+// Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
+// weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
+// first along the reduction axis, two codes per stored word.
+const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const FC1_PACKED_COLS: u32 = {{ fc1PackedCols }}u;
+{% if not colWiseScales %}
+// Block-wise scales (fc1_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
+// axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
+const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const FC1_BLOCKS: u32 = {{ fc1ScaleBlocks }}u;
+{% endif %}
+fn fc1_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * FC1_ROWS + row;
+{% if packSize == 2 %}
+  let word = fc1_experts_weights[base * FC1_PACKED_COLS + col / 2u];
+  let code = f32((word >> ((col & 1u) * 4u)) & 15u);
+{% else %}
+  let code = f32(fc1_experts_weights[base * FC1_PACKED_COLS + col] & 255u);
+{% endif %}
+{% if not colWiseScales %}
+  return (code - MIDPOINT) * fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK];
+{% else %}
+  return (code - MIDPOINT) * fc1_scales[base];
+{% endif %}
+}
+// 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 %}
+  let pair = base * FC1_PACKED_COLS + col / 2u;
+  let w0 = fc1_experts_weights[pair];
+  let w1 = fc1_experts_weights[pair + 1u];
+  let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
+                        f32(w1 & 15u), f32((w1 >> 4u) & 15u));
+{% else %}
+  let word = base * FC1_PACKED_COLS + col;
+  let codes = vec4<f32>(f32(fc1_experts_weights[word] & 255u), f32(fc1_experts_weights[word + 1u] & 255u),
+                        f32(fc1_experts_weights[word + 2u] & 255u), f32(fc1_experts_weights[word + 3u] & 255u));
+{% endif %}
+{% if not colWiseScales %}
+  let scale = fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK];
+{% else %}
+  let scale = fc1_scales[base];
+{% endif %}
+  return (codes - vec4<f32>(MIDPOINT)) * scale;
+}
+{% if activationType == "swiglu" %}
+fn activate(gate_in: f32, up_in: f32) -> f32 {
+{% if hasSwigluLimit %}
+  // swiglu_limit clamps gate from above and up to [-limit, limit] before the
+  // product; when absent, no clamp is applied.
+  let gate = min(gate_in, {{ swigluLimit }});
+  let up = clamp(up_in, -({{ swigluLimit }}), {{ swigluLimit }});
+{% else %}
+  let gate = gate_in;
+  let up = up_in;
+{% endif %}
+  return gate / (1.0 + exp(-params.activationAlpha * gate)) * (up + params.activationBeta);
+}
+{% else %}
+fn activate(value: f32) -> f32 {
+  return max(value, 0.0);
+}
+{% endif %}
+const KVEC: u32 = {{ groupTileKVec }}u;
+{% macro stage_group_tiles(aLoad, bLoad, b2Load, kExtent, nExtent, guarded, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
+      for (var idx = tid; idx < MTILE * KVEC; idx = idx + THREADS) {
+        let m = idx / KVEC;
+        let kb = (idx % KVEC) * 4u;
+        var a_vec: vec4<f32>;
+{% if aLoadVec4 and not guarded %}
+        a_vec = {{ aLoadVec4 }};
+{% else %}
+{% for j in range(4) %}
+        {
+          let k = kb + {{ j }}u;
+          a_vec[{{ j }}] = {{ ("select(0.0, " ~ aLoad ~ ", k_base + k < " ~ kExtent ~ ")") if guarded else aLoad }};
+        }
+{% endfor %}
+{% endif %}
+        a_tile[idx] = a_vec;
+      }
+      for (var idx = tid; idx < NTILE * KVEC; idx = idx + THREADS) {
+        let n = idx / KVEC;
+        let kb = (idx % KVEC) * 4u;
+        let col = n_base + n;
+{% if guarded %}
+        let col_live = col < {{ nExtent }};
+{% endif %}
+        var b_vec: vec4<f32>;
+{% if b2Load %}
+        var b2_vec: vec4<f32>;
+{% endif %}
+{% if bLoadVec4 and not guarded %}
+        b_vec = {{ bLoadVec4 }};
+{% if b2Load %}
+        b2_vec = {{ b2LoadVec4 }};
+{% endif %}
+{% else %}
+{% for j in range(4) %}
+        {
+          let k = kb + {{ j }}u;
+{% if guarded %}
+          let live = col_live && k_base + k < {{ kExtent }};
+{% endif %}
+          b_vec[{{ j }}] = {{ ("select(0.0, " ~ bLoad ~ ", live)") if guarded else bLoad }};
+{% if b2Load %}
+          b2_vec[{{ j }}] = {{ ("select(0.0, " ~ b2Load ~ ", live)") if guarded else b2Load }};
+{% endif %}
+        }
+{% endfor %}
+{% endif %}
+        b_tile[idx] = b_vec;
+{% if b2Load %}
+        b2_tile[idx] = b2_vec;
+{% endif %}
+      }
+{%- endmacro %}
+{% macro group_tile_loop(aLoad, bLoad, b2Load, kExtent, nExtent, regM, regN, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
+{% for r in range(regM) %}
+{% for c in range(regN) %}
+  var acc_{{ r }}_{{ c }} = 0.0;
+{% if b2Load %}
+  var acc2_{{ r }}_{{ c }} = 0.0;
+{% endif %}
+{% endfor %}
+{% endfor %}
+  // A tile that lies wholly inside both extents needs no per-element bounds check, and
+  // that is every tile but the edges: the guards below exist for a reduction axis or an
+  // output width the tile size does not divide, yet the interior tiles were paying four
+  // comparisons and four selects per staged vec4 -- against 64 fused multiply-adds of
+  // actual work. Both arms stage identical values. The test is workgroup-uniform, so the
+  // barrier after it is still reached by every invocation together.
+  let n_full = n_base + NTILE <= {{ nExtent }};
+  var k_base = 0u;
+  loop {
+    if (k_base >= {{ kExtent }}) {
+      break;
+    }
+    if (n_full && k_base + KTILE <= {{ kExtent }}) {
+{{ stage_group_tiles(aLoad, bLoad, b2Load, kExtent, nExtent, false, bLoadVec4, b2LoadVec4, aLoadVec4) }}
+    } else {
+{{ stage_group_tiles(aLoad, bLoad, b2Load, kExtent, nExtent, true) }}
+    }
+    workgroupBarrier();
+    for (var kv = 0u; kv < KVEC; kv = kv + 1u) {
+{% for r in range(regM) %}
+      let a{{ r }} = a_tile[(lid.y * {{ regM }}u + {{ r }}u) * KVEC + kv];
+{% endfor %}
+{% for c in range(regN) %}
+      let b{{ c }} = b_tile[(lid.x * {{ regN }}u + {{ c }}u) * KVEC + kv];
+{% if b2Load %}
+      let s{{ c }} = b2_tile[(lid.x * {{ regN }}u + {{ c }}u) * KVEC + kv];
+{% endif %}
+{% endfor %}
+{% for r in range(regM) %}
+{% for c in range(regN) %}
+      acc_{{ r }}_{{ c }} = acc_{{ r }}_{{ c }} + dot(a{{ r }}, b{{ c }});
+{% if b2Load %}
+      acc2_{{ r }}_{{ c }} = acc2_{{ r }}_{{ c }} + dot(a{{ r }}, s{{ c }});
+{% endif %}
+{% endfor %}
+{% endfor %}
+    }
+    // Orders this step's tile reads before the next step overwrites them.
+    workgroupBarrier();
+    k_base = k_base + KTILE;
+  }
+{%- endmacro %}
+var<workgroup> row_slot: array<u32, {{ groupTileM }}>;
+var<workgroup> a_tile: array<vec4<f32>, {{ groupTileM * groupTileKVec }}>;
+var<workgroup> b_tile: array<vec4<f32>, {{ groupTileN * groupTileKVec }}>;
+{% if isSwiglu %}
+var<workgroup> b2_tile: array<vec4<f32>, {{ groupTileN * groupTileKVec }}>;
+{% endif %}
+@compute @workgroup_size(THREADS_SIDE, THREADS_SIDE, 1)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  let tile = wid.x;
+  if (tile >= tile_meta[0]) {
+    return;
+  }
+  let expert = tile_meta[1u + tile * 3u];
+  let slice_base = tile_meta[1u + tile * 3u + 1u];
+  let rows = tile_meta[1u + tile * 3u + 2u];
+  let n_base = wid.y * NTILE;
+  let tid = lid.y * THREADS_SIDE + lid.x;
+  for (var m = tid; m < MTILE; m = m + THREADS) {
+    // Rows past this tile's end park on row 0; their results are discarded at the store.
+    row_slot[m] = select(0u, slot_list[slice_base + m], m < rows);
+  }
+  workgroupBarrier();
+{{ group_tile_loop("f32(input[(params.tokenOffset + row_slot[m] / TOP_K) * HIDDEN + k_base + k])",
+                   "fc1_weight(expert, " ~ primaryRow ~ ", k_base + k)",
+                   "fc1_weight(expert, col * 2u + 1u, k_base + k)" if isSwiglu else "",
+                   "HIDDEN", "INTER", regM, regN,
+                   "fc1_weight_vec4(expert, " ~ primaryRow ~ ", k_base + kb)",
+                   "fc1_weight_vec4(expert, col * 2u + 1u, k_base + kb)" if isSwiglu else "") }}
+{% for r in range(regM) %}
+{% for c in range(regN) %}
+  {
+    let m = lid.y * {{ regM }}u + {{ r }}u;
+    let col = n_base + lid.x * {{ regN }}u + {{ c }}u;
+    if (m < rows && col < INTER) {
+{% if isSwiglu %}
+      hidden_act[row_slot[m] * INTER + col] = activate(acc_{{ r }}_{{ c }}, acc2_{{ r }}_{{ c }});
+{% else %}
+      hidden_act[row_slot[m] * INTER + col] = activate(acc_{{ r }}_{{ c }});
+{% endif %}
+    }
+  }
+{% endfor %}
+{% endfor %}
+}

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

	@@ -0,0 +1,116 @@

+{{ env.wgsl.resourceDeclarations }}
+const HIDDEN: u32 = {{ hidden }}u;
+const FC1_ROWS: u32 = {{ fc1Rows }}u;
+const INTER: u32 = {{ inter }}u;
+// WGSL rejects a const-expression `/ 0u` outright; guard the divisor. INTER == 0 makes `total` 0,
+// so the index bounds-guard returns every thread before this divisor is used (degenerate empty
+// inter-dim case).
+const INTER_DIV: u32 = max(1u, INTER);
+const WG: u32 = {{ workgroupSize }}u;
+{% if topK > 1 %}
+const TOP_K: u32 = {{ topK }}u;
+{% endif %}
+// Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
+// weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
+// first along the reduction axis, two codes per stored word.
+const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const FC1_PACKED_COLS: u32 = {{ fc1PackedCols }}u;
+{% if not colWiseScales %}
+// Block-wise scales (fc1_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
+// axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
+const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const FC1_BLOCKS: u32 = {{ fc1ScaleBlocks }}u;
+{% endif %}
+fn fc1_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * FC1_ROWS + row;
+{% if packSize == 2 %}
+  let word = fc1_experts_weights[base * FC1_PACKED_COLS + col / 2u];
+  let code = f32((word >> ((col & 1u) * 4u)) & 15u);
+{% else %}
+  let code = f32(fc1_experts_weights[base * FC1_PACKED_COLS + col] & 255u);
+{% endif %}
+{% if not colWiseScales %}
+  return (code - MIDPOINT) * fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK];
+{% else %}
+  return (code - MIDPOINT) * fc1_scales[base];
+{% endif %}
+}
+{% if activationType == "swiglu" %}
+fn activate(gate_in: f32, up_in: f32) -> f32 {
+{% if hasSwigluLimit %}
+  // swiglu_limit clamps gate from above and up to [-limit, limit] before the
+  // product; when absent, no clamp is applied.
+  let gate = min(gate_in, {{ swigluLimit }});
+  let up = clamp(up_in, -({{ swigluLimit }}), {{ swigluLimit }});
+{% else %}
+  let gate = gate_in;
+  let up = up_in;
+{% endif %}
+  return gate / (1.0 + exp(-params.activationAlpha * gate)) * (up + params.activationBeta);
+}
+{% else %}
+fn activate(value: f32) -> f32 {
+  return max(value, 0.0);
+}
+{% 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 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.
+  let total = params.tokenCount * TOP_K * INTER;
+  if (index >= total) {
+    return;
+  }
+  let slot = index / INTER_DIV;
+  let local_token = slot / TOP_K;
+  let token = params.tokenOffset + local_token;
+  let i = index % INTER_DIV;
+  let expert = route_expert[token * TOP_K + slot % TOP_K];
+{% else %}
+  let total = params.tokenCount * INTER;
+  if (index >= total) {
+    return;
+  }
+  let local_token = index / INTER_DIV;
+  let token = params.tokenOffset + local_token;
+  let i = index % INTER_DIV;
+  let expert = route_expert[token];
+{% endif %}
+  var gate = 0.0;
+{% if activationType == "swiglu" %}
+  var up = 0.0;
+{% endif %}
+  for (var h: u32 = 0u; h < HIDDEN; h = h + 1u) {
+    let x = f32(input[token * HIDDEN + h]);
+{% if activationType == "swiglu" %}
+    gate = gate + x * fc1_weight(expert, i * 2u, h);
+{% else %}
+    gate = gate + x * fc1_weight(expert, i, h);
+{% endif %}
+{% if activationType == "swiglu" %}
+    up = up + x * fc1_weight(expert, i * 2u + 1u, h);
+{% endif %}
+  }
+{% if activationType == "swiglu" %}
+  hidden_act[index] = activate(gate, up);
+{% else %}
+  hidden_act[index] = activate(gate);
+{% endif %}
+}

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

	@@ -0,0 +1,33 @@

+{{ env.wgsl.resourceDeclarations }}
+// Routed sum for the grouped schedule. The grouped FC2 stage produces one projected row per
+// routed SLOT because a token's slots are spread across different expert tiles; this pass
+// applies each slot's routing weight and collapses the TOP_K slots into the token's output row.
+// The expert's FC2 bias is already inside its slot's row, so it is scaled by that slot's weight
+// exactly as the other schedules scale it.
+const HIDDEN: u32 = {{ hidden }}u;
+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;
+  }
+  let local_token = index / HIDDEN;
+  let token = params.tokenOffset + local_token;
+  let out_col = index % HIDDEN;
+  var acc = 0.0;
+  for (var j = 0u; j < TOP_K; j = j + 1u) {
+    acc = acc + slot_out[(local_token * TOP_K + j) * HIDDEN + out_col] * route_mix[token * TOP_K + j];
+  }
+  output[token * HIDDEN + out_col] = acc;
+}

build/webgpu/qmoe-output-gemv.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,117 @@

+{{ env.wgsl.resourceDeclarations }}
+const HIDDEN: u32 = {{ hidden }}u;
+const INTER: u32 = {{ inter }}u;
+const TOP_K: u32 = {{ topK }}u;
+// A workgroup owns ROWS output columns of one token and LANES lanes split the
+// intermediate axis, so the split stage's per-thread INTER-long walk becomes a
+// LANES-wide cooperative reduction.
+const LANES: u32 = {{ decodeLanes }}u;
+const ROWS: u32 = {{ decodeRows }}u;
+// Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
+// weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
+// first along the reduction axis, two codes per stored word.
+const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const FC2_PACKED_COLS: u32 = {{ fc2PackedCols }}u;
+{% if not colWiseScales %}
+// Block-wise scales (fc2_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
+// axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
+const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const FC2_BLOCKS: u32 = {{ fc2ScaleBlocks }}u;
+{% endif %}
+fn fc2_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * HIDDEN + row;
+{% if packSize == 2 %}
+  let word = fc2_experts_weights[base * FC2_PACKED_COLS + col / 2u];
+  let code = f32((word >> ((col & 1u) * 4u)) & 15u);
+{% else %}
+  let code = f32(fc2_experts_weights[base * FC2_PACKED_COLS + col] & 255u);
+{% endif %}
+{% if not colWiseScales %}
+  return (code - MIDPOINT) * fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK];
+{% else %}
+  return (code - MIDPOINT) * fc2_scales[base];
+{% endif %}
+}
+// Decode two consecutive reduction-axis codes together. Four-bit codes share a
+// stored word; eight-bit codes occupy adjacent words.
+fn fc2_weight_pair(expert: u32, row: u32, col: u32) -> vec2<f32> {
+  let base = expert * HIDDEN + row;
+{% if packSize == 2 %}
+  let word = fc2_experts_weights[base * FC2_PACKED_COLS + col / 2u];
+  let lo = f32(word & 15u) - MIDPOINT;
+  let hi = f32((word >> 4u) & 15u) - MIDPOINT;
+{% else %}
+  let lo = f32(fc2_experts_weights[base * FC2_PACKED_COLS + col] & 255u) - MIDPOINT;
+  let hi = f32(fc2_experts_weights[base * FC2_PACKED_COLS + col + 1u] & 255u) - MIDPOINT;
+{% endif %}
+{% if not colWiseScales %}
+  return vec2<f32>(lo * fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK],
+                   hi * fc2_scales[base * FC2_BLOCKS + (col + 1u) / QUANT_BLOCK]);
+{% else %}
+  let scale = fc2_scales[base];
+  return vec2<f32>(lo * scale, hi * scale);
+{% endif %}
+}
+{% macro rowlane_fold(arrays, lanes="LANES", lane="lane", slot="slot") %}
+  var n = {{ lanes }};
+  while (n > 1u) {
+    let half = (n + 1u) / 2u;
+    if ({{ lane }} + half < n) {
+{% for a in arrays %}
+      {{ a }}[{{ slot }}] = {{ a }}[{{ slot }}] + {{ a }}[{{ slot }} + half];
+{% endfor %}
+    }
+    workgroupBarrier();
+    n = half;
+  }
+{%- endmacro %}
+var<workgroup> partial: array<f32, {{ decodeLanes * decodeRows }}>;
+@compute @workgroup_size(LANES, ROWS, 1)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  let out_col = wid.x * ROWS + lid.y;
+  let lane = lid.x;
+  let local_token = wid.y;
+  let token = params.tokenOffset + local_token;
+  let in_range = out_col < HIDDEN;
+  var acc = 0.0;
+  if (in_range) {
+    // route_mix is uniform across a workgroup's lanes, so folding each slot's
+    // share into its own partial here is the same value as scaling the reduced
+    // sum -- and it leaves ONE cross-lane fold covering all TOP_K slots.
+    for (var j = 0u; j < TOP_K; j = j + 1u) {
+      let route_slot = token * TOP_K + j;
+      let expert = route_expert[route_slot];
+      let act_base = (local_token * TOP_K + j) * INTER;
+      var part = 0.0;
+      // Two reduction columns per step: at 4 bits they are one stored word.
+      for (var c = lane * 2u; c + 1u < INTER; c = c + LANES * 2u) {
+        let p = fc2_weight_pair(expert, out_col, c);
+        part = part + hidden_act[act_base + c] * p.x
+                    + hidden_act[act_base + c + 1u] * p.y;
+      }
+      // An odd reduction length leaves one column, which only 8-bit weights can have.
+      if ((INTER & 1u) == 1u && lane == 0u) {
+        let i = INTER - 1u;
+        part = part + hidden_act[act_base + i] * fc2_weight(expert, out_col, i);
+      }
+      acc = acc + part * route_mix[route_slot];
+    }
+  }
+  let slot = lid.y * LANES + lane;
+  partial[slot] = acc;
+  workgroupBarrier();
+{{ rowlane_fold(["partial"]) }}
+  if (lane == 0u && in_range) {
+    output[token * HIDDEN + out_col] = partial[lid.y * LANES];
+  }
+}

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

	@@ -0,0 +1,196 @@

+enable subgroups;
+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;
+const SUB_ROWS: u32 = 16u;
+// Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
+// weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
+// first along the reduction axis, two codes per stored word.
+const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const FC2_PACKED_COLS: u32 = {{ fc2PackedCols }}u;
+{% if not colWiseScales %}
+// Block-wise scales (fc2_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
+// axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
+const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const FC2_BLOCKS: u32 = {{ fc2ScaleBlocks }}u;
+{% endif %}
+// 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 %}
+  let pair = base * FC2_PACKED_COLS + col / 2u;
+  let w0 = fc2_experts_weights[pair];
+  let w1 = fc2_experts_weights[pair + 1u];
+  let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
+                        f32(w1 & 15u), f32((w1 >> 4u) & 15u));
+{% else %}
+  let word = base * FC2_PACKED_COLS + col;
+  let codes = vec4<f32>(f32(fc2_experts_weights[word] & 255u), f32(fc2_experts_weights[word + 1u] & 255u),
+                        f32(fc2_experts_weights[word + 2u] & 255u), f32(fc2_experts_weights[word + 3u] & 255u));
+{% endif %}
+{% if not colWiseScales %}
+  let scale = fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK];
+{% else %}
+  let scale = fc2_scales[base];
+{% endif %}
+  return (codes - vec4<f32>(MIDPOINT)) * scale;
+}
+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];
+    }
+  } else {
+    for (var i = 0u; i < 8u; i = i + 1u) {
+      tile_A[row * TILE_K + col + i] = 0.0;
+    }
+  }
+}
+// Columns past HIDDEN stage zero: they contribute nothing to the product and are never
+// stored, and their packed reads would be out of range.
+fn loadSHMB(expert: u32, col_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+  let n = col_base + row;
+  let col = c_idx * 16u;
+  for (var i = 0u; i < 16u; i = i + 4u) {
+    let k = k_idx + col + i;
+    var quad = vec4<f32>(0.0);
+    if (n < HIDDEN) {
+      quad = fc2_weight_vec4(expert, n, k);
+    }
+    tile_B[row * TILE_K + col + i] = quad.x;
+    tile_B[row * TILE_K + col + i + 1u] = quad.y;
+    tile_B[row * TILE_K + col + i + 2u] = quad.z;
+    tile_B[row * TILE_K + col + i + 3u] = quad.w;
+  }
+}
+// One output column of an 8-wide result block, scattered back to its routed slot.
+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);
+{% endfor %}
+  }
+}
+@compute @workgroup_size({{ sgmatWorkgroup }}, 1, 1)
+fn main(
+  @builtin(workgroup_id) wid: vec3<u32>,
+  @builtin(local_invocation_index) local_idx: u32,
+  @builtin(subgroup_invocation_id) sg_id: u32,
+  @builtin(subgroup_size) sg_size: u32
+) {
+  let tile = wid.x;
+  if (tile >= tile_meta[0]) {
+    return;
+  }
+  let expert = tile_meta[1u + tile * 3u];
+  let slice_base = tile_meta[1u + tile * 3u + 1u];
+  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;
+  let subtile_idy = subtile_id % {{ sgmatRowSubtiles }}u;
+  let base_A = subtile_idy * SUB_ROWS;
+  let base_B = subtile_idx * SUB_COLS;
+{% for row in [0, 1] %}{% for col in [0, 1, 2, 3] %}
+  var matC{{ row }}{{ col }}: subgroup_matrix_result<f32, 8, 8>;
+{% 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>>(&tile_A, matrix_a_offset, false, TILE_K);
+      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, 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>>(&tile_B, matrix_b_offset, true, TILE_K);
+      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
+      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
+      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
+      matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
+      matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
+      matC02 = subgroupMatrixMultiplyAccumulate(matA0, matB2, matC02);
+      matC03 = subgroupMatrixMultiplyAccumulate(matA0, matB3, matC03);
+      matC10 = subgroupMatrixMultiplyAccumulate(matA1, matB0, matC10);
+      matC11 = subgroupMatrixMultiplyAccumulate(matA1, matB1, matC11);
+      matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12);
+      matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13);
+    }
+    workgroupBarrier();
+  }
+  // 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(&scratch[subtile_id][{{ col }}], 0u, matC0{{ col }}, false, 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(&scratch[subtile_id][{{ col }}], 0u, matC1{{ col }}, false, 8u);
+{% endfor %}
+  workgroupBarrier();
+  storeOutput(col_base + base_B, row, col, subtile_id, base_A + 8u, rows);
+}

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

	@@ -0,0 +1,214 @@

+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.QMoE, expert-grouped FC2 stage, sharing the tile loop with the grouped FC1
+// stage: the reduction is the intermediate axis and the output width is the hidden axis.
+//
+// The result is written per slot rather than per token because routed slots land in
+// different expert tiles and nothing here can sum them. The mix stage that follows applies each
+// slot's routing weight and collapses the TOP_K slots into the token's output row.
+const HIDDEN: u32 = {{ hidden }}u;
+const INTER: u32 = {{ inter }}u;
+const MTILE: u32 = {{ groupTileM }}u;
+const NTILE: u32 = {{ groupTileN }}u;
+const KTILE: u32 = {{ groupTileK }}u;
+const THREADS_SIDE: u32 = {{ groupThreads }}u;
+const THREADS: u32 = THREADS_SIDE * THREADS_SIDE;
+// Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
+// weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
+// first along the reduction axis, two codes per stored word.
+const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const FC2_PACKED_COLS: u32 = {{ fc2PackedCols }}u;
+{% if not colWiseScales %}
+// Block-wise scales (fc2_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
+// axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
+const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const FC2_BLOCKS: u32 = {{ fc2ScaleBlocks }}u;
+{% endif %}
+fn fc2_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * HIDDEN + row;
+{% if packSize == 2 %}
+  let word = fc2_experts_weights[base * FC2_PACKED_COLS + col / 2u];
+  let code = f32((word >> ((col & 1u) * 4u)) & 15u);
+{% else %}
+  let code = f32(fc2_experts_weights[base * FC2_PACKED_COLS + col] & 255u);
+{% endif %}
+{% if not colWiseScales %}
+  return (code - MIDPOINT) * fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK];
+{% else %}
+  return (code - MIDPOINT) * fc2_scales[base];
+{% endif %}
+}
+// 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 %}
+  let pair = base * FC2_PACKED_COLS + col / 2u;
+  let w0 = fc2_experts_weights[pair];
+  let w1 = fc2_experts_weights[pair + 1u];
+  let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
+                        f32(w1 & 15u), f32((w1 >> 4u) & 15u));
+{% else %}
+  let word = base * FC2_PACKED_COLS + col;
+  let codes = vec4<f32>(f32(fc2_experts_weights[word] & 255u), f32(fc2_experts_weights[word + 1u] & 255u),
+                        f32(fc2_experts_weights[word + 2u] & 255u), f32(fc2_experts_weights[word + 3u] & 255u));
+{% endif %}
+{% if not colWiseScales %}
+  let scale = fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK];
+{% else %}
+  let scale = fc2_scales[base];
+{% endif %}
+  return (codes - vec4<f32>(MIDPOINT)) * scale;
+}
+const KVEC: u32 = {{ groupTileKVec }}u;
+{% macro stage_group_tiles(aLoad, bLoad, b2Load, kExtent, nExtent, guarded, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
+      for (var idx = tid; idx < MTILE * KVEC; idx = idx + THREADS) {
+        let m = idx / KVEC;
+        let kb = (idx % KVEC) * 4u;
+        var a_vec: vec4<f32>;
+{% if aLoadVec4 and not guarded %}
+        a_vec = {{ aLoadVec4 }};
+{% else %}
+{% for j in range(4) %}
+        {
+          let k = kb + {{ j }}u;
+          a_vec[{{ j }}] = {{ ("select(0.0, " ~ aLoad ~ ", k_base + k < " ~ kExtent ~ ")") if guarded else aLoad }};
+        }
+{% endfor %}
+{% endif %}
+        a_tile[idx] = a_vec;
+      }
+      for (var idx = tid; idx < NTILE * KVEC; idx = idx + THREADS) {
+        let n = idx / KVEC;
+        let kb = (idx % KVEC) * 4u;
+        let col = n_base + n;
+{% if guarded %}
+        let col_live = col < {{ nExtent }};
+{% endif %}
+        var b_vec: vec4<f32>;
+{% if b2Load %}
+        var b2_vec: vec4<f32>;
+{% endif %}
+{% if bLoadVec4 and not guarded %}
+        b_vec = {{ bLoadVec4 }};
+{% if b2Load %}
+        b2_vec = {{ b2LoadVec4 }};
+{% endif %}
+{% else %}
+{% for j in range(4) %}
+        {
+          let k = kb + {{ j }}u;
+{% if guarded %}
+          let live = col_live && k_base + k < {{ kExtent }};
+{% endif %}
+          b_vec[{{ j }}] = {{ ("select(0.0, " ~ bLoad ~ ", live)") if guarded else bLoad }};
+{% if b2Load %}
+          b2_vec[{{ j }}] = {{ ("select(0.0, " ~ b2Load ~ ", live)") if guarded else b2Load }};
+{% endif %}
+        }
+{% endfor %}
+{% endif %}
+        b_tile[idx] = b_vec;
+{% if b2Load %}
+        b2_tile[idx] = b2_vec;
+{% endif %}
+      }
+{%- endmacro %}
+{% macro group_tile_loop(aLoad, bLoad, b2Load, kExtent, nExtent, regM, regN, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
+{% for r in range(regM) %}
+{% for c in range(regN) %}
+  var acc_{{ r }}_{{ c }} = 0.0;
+{% if b2Load %}
+  var acc2_{{ r }}_{{ c }} = 0.0;
+{% endif %}
+{% endfor %}
+{% endfor %}
+  // A tile that lies wholly inside both extents needs no per-element bounds check, and
+  // that is every tile but the edges: the guards below exist for a reduction axis or an
+  // output width the tile size does not divide, yet the interior tiles were paying four
+  // comparisons and four selects per staged vec4 -- against 64 fused multiply-adds of
+  // actual work. Both arms stage identical values. The test is workgroup-uniform, so the
+  // barrier after it is still reached by every invocation together.
+  let n_full = n_base + NTILE <= {{ nExtent }};
+  var k_base = 0u;
+  loop {
+    if (k_base >= {{ kExtent }}) {
+      break;
+    }
+    if (n_full && k_base + KTILE <= {{ kExtent }}) {
+{{ stage_group_tiles(aLoad, bLoad, b2Load, kExtent, nExtent, false, bLoadVec4, b2LoadVec4, aLoadVec4) }}
+    } else {
+{{ stage_group_tiles(aLoad, bLoad, b2Load, kExtent, nExtent, true) }}
+    }
+    workgroupBarrier();
+    for (var kv = 0u; kv < KVEC; kv = kv + 1u) {
+{% for r in range(regM) %}
+      let a{{ r }} = a_tile[(lid.y * {{ regM }}u + {{ r }}u) * KVEC + kv];
+{% endfor %}
+{% for c in range(regN) %}
+      let b{{ c }} = b_tile[(lid.x * {{ regN }}u + {{ c }}u) * KVEC + kv];
+{% if b2Load %}
+      let s{{ c }} = b2_tile[(lid.x * {{ regN }}u + {{ c }}u) * KVEC + kv];
+{% endif %}
+{% endfor %}
+{% for r in range(regM) %}
+{% for c in range(regN) %}
+      acc_{{ r }}_{{ c }} = acc_{{ r }}_{{ c }} + dot(a{{ r }}, b{{ c }});
+{% if b2Load %}
+      acc2_{{ r }}_{{ c }} = acc2_{{ r }}_{{ c }} + dot(a{{ r }}, s{{ c }});
+{% endif %}
+{% endfor %}
+{% endfor %}
+    }
+    // Orders this step's tile reads before the next step overwrites them.
+    workgroupBarrier();
+    k_base = k_base + KTILE;
+  }
+{%- endmacro %}
+var<workgroup> row_slot: array<u32, {{ groupTileM }}>;
+var<workgroup> a_tile: array<vec4<f32>, {{ groupTileM * groupTileKVec }}>;
+var<workgroup> b_tile: array<vec4<f32>, {{ groupTileN * groupTileKVec }}>;
+@compute @workgroup_size(THREADS_SIDE, THREADS_SIDE, 1)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  let tile = wid.x;
+  if (tile >= tile_meta[0]) {
+    return;
+  }
+  let expert = tile_meta[1u + tile * 3u];
+  let slice_base = tile_meta[1u + tile * 3u + 1u];
+  let rows = tile_meta[1u + tile * 3u + 2u];
+  let n_base = wid.y * NTILE;
+  let tid = lid.y * THREADS_SIDE + lid.x;
+  for (var m = tid; m < MTILE; m = m + THREADS) {
+    // Rows past this tile's end park on row 0; their results are discarded at the store.
+    row_slot[m] = select(0u, slot_list[slice_base + m], m < rows);
+  }
+  workgroupBarrier();
+{{ group_tile_loop("hidden_act[row_slot[m] * INTER + k_base + k]",
+                   "fc2_weight(expert, col, k_base + k)",
+                   "", "INTER", "HIDDEN", regM, regN,
+                   "fc2_weight_vec4(expert, col, k_base + kb)") }}
+{% for r in range(regM) %}
+{% for c in range(regN) %}
+  {
+    let m = lid.y * {{ regM }}u + {{ r }}u;
+    let col = n_base + lid.x * {{ regN }}u + {{ c }}u;
+    if (m < rows && col < HIDDEN) {
+      slot_out[row_slot[m] * HIDDEN + col] = acc_{{ r }}_{{ c }};
+    }
+  }
+{% endfor %}
+{% endfor %}
+}

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

	@@ -0,0 +1,78 @@

+{{ env.wgsl.resourceDeclarations }}
+const HIDDEN: u32 = {{ hidden }}u;
+const INTER: u32 = {{ inter }}u;
+const WG: u32 = {{ workgroupSize }}u;
+{% if topK > 1 %}
+const TOP_K: u32 = {{ topK }}u;
+{% endif %}
+// Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
+// weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
+// first along the reduction axis, two codes per stored word.
+const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const FC2_PACKED_COLS: u32 = {{ fc2PackedCols }}u;
+{% if not colWiseScales %}
+// Block-wise scales (fc2_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
+// axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
+const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const FC2_BLOCKS: u32 = {{ fc2ScaleBlocks }}u;
+{% endif %}
+fn fc2_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * HIDDEN + row;
+{% if packSize == 2 %}
+  let word = fc2_experts_weights[base * FC2_PACKED_COLS + col / 2u];
+  let code = f32((word >> ((col & 1u) * 4u)) & 15u);
+{% else %}
+  let code = f32(fc2_experts_weights[base * FC2_PACKED_COLS + col] & 255u);
+{% endif %}
+{% if not colWiseScales %}
+  return (code - MIDPOINT) * fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK];
+{% else %}
+  return (code - MIDPOINT) * fc2_scales[base];
+{% 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 index = gid.x + gid.y * nwg.x * WG;
+  let total = params.tokenCount * HIDDEN;
+  if (index >= total) {
+    return;
+  }
+  let local_token = index / HIDDEN;
+  let token = params.tokenOffset + local_token;
+  let out_col = index % HIDDEN;
+{% if topK > 1 %}
+  // Each routed slot contributes its own fc2 projection of its own activation row, weighted by
+  // that slot's share. Summing per slot (rather than mixing the activations first) keeps the
+  // K == 1 term exactly the `acc * mix` the top-1 stage computes.
+  var acc = 0.0;
+  for (var j = 0u; j < TOP_K; j = j + 1u) {
+    let expert = route_expert[token * TOP_K + j];
+    var part = 0.0;
+    for (var i: u32 = 0u; i < INTER; i = i + 1u) {
+      part = part + hidden_act[(local_token * TOP_K + j) * INTER + i] * fc2_weight(expert, out_col, i);
+    }
+    acc = acc + part * route_mix[token * TOP_K + j];
+  }
+  output[token * HIDDEN + out_col] = acc;
+{% else %}
+  let expert = route_expert[token];
+  let mix = route_mix[token];
+  var acc = 0.0;
+  for (var i: u32 = 0u; i < INTER; i = i + 1u) {
+    acc = acc + hidden_act[local_token * INTER + i] * fc2_weight(expert, out_col, i);
+  }
+  output[token * HIDDEN + out_col] = acc * mix;
+{% endif %}
+}

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

	@@ -0,0 +1,12 @@

+{{ env.wgsl.resourceDeclarations }}
+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;
+  }
+}

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

	@@ -0,0 +1,105 @@

+{{ env.wgsl.resourceDeclarations }}
+const TOKENS: u32 = {{ tokens }}u;
+const EXPERTS: u32 = {{ experts }}u;
+const WG: u32 = {{ workgroupSize }}u;
+{% if topK > 1 %}
+const TOP_K: u32 = {{ topK }}u;
+// Serial top-K selection in O(TOP_K * EXPERTS), followed by a stable softmax
+// over the selected logits. Equal logits select the higher expert index.
+struct TopKRoute {
+  experts: array<u32, TOP_K>,
+  weights: array<f32, TOP_K>,
+}
+fn topk_route(token: u32) -> TopKRoute {
+  var result: TopKRoute;
+  for (var j = 0u; j < TOP_K; j = j + 1u) {
+    // EXPERTS is the "nothing picked yet" sentinel; the contract keeps TOP_K <= EXPERTS, so
+    // every pass finds an untaken expert and no slot can keep it.
+    var best_expert = EXPERTS;
+    var best_weight = 0.0;
+    for (var e = 0u; e < EXPERTS; e = e + 1u) {
+      var taken = false;
+      for (var p = 0u; p < j; p = p + 1u) {
+        if (result.experts[p] == e) {
+          taken = true;
+        }
+      }
+      if (!taken) {
+        let candidate = router_probs[token * EXPERTS + e];
+        if (best_expert == EXPERTS || candidate > best_weight ||
+            (candidate == best_weight && e > best_expert)) {
+          best_expert = e;
+          best_weight = candidate;
+        }
+      }
+    }
+    result.experts[j] = best_expert;
+    result.weights[j] = best_weight;
+  }
+  // QMoE without a separate router_weights input always applies softmax over
+  // the selected top-k logits, independently of normalize_routing_weights.
+  let max_logit = result.weights[0];
+  var sum = 0.0;
+  for (var j = 0u; j < TOP_K; j = j + 1u) {
+    result.weights[j] = exp(result.weights[j] - max_logit);
+    sum = sum + result.weights[j];
+  }
+  if (sum != 0.0) {
+    for (var j = 0u; j < TOP_K; j = j + 1u) {
+      result.weights[j] = result.weights[j] / sum;
+    }
+  }
+  return result;
+}
+{% else %}
+// Serial top-1 selection over `(logit, expert_index)` pairs ordered descending,
+// so equal logits select the higher expert index.
+struct Top1Route {
+  expert: u32,
+  weight: f32,
+}
+fn top1_route(token: u32) -> Top1Route {
+  var result = Top1Route(0u, 1.0);
+  var best_weight = router_probs[token * EXPERTS];
+  for (var e = 1u; e < EXPERTS; e += 1u) {
+    let candidate = router_probs[token * EXPERTS + e];
+    if (candidate >= best_weight) {
+      result.expert = e;
+      best_weight = candidate;
+    }
+  }
+  return result;
+}
+{% 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;
+  }
+{% if topK > 1 %}
+  // Slot j of each token's TOP_K-wide run holds its j-th ranked expert and
+  // that expert's softmax weight within the selected set.
+  let route = topk_route(token);
+  for (var j = 0u; j < TOP_K; j = j + 1u) {
+    route_expert[token * TOP_K + j] = route.experts[j];
+    route_mix[token * TOP_K + j] = route.weights[j];
+  }
+{% else %}
+  let route = top1_route(token);
+  route_expert[token] = route.expert;
+  route_mix[token] = route.weight;
+{% endif %}
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,622 @@

+{
+  "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": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [1.0, -2.0, 0.5, 3.0, -1.0, 2.0, 4.0, -0.5] }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [0.1, 1.2, -0.4, 2.0, -1.0, 0.5] }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [3, 4, 2],
+          "data": { "kind": "cycle", "values": [154, 104, 124, 185, 39, 211, 78] }
+        },
+        "fc1ScalesT": { "dtype": "float32", "shape": [3, 4], "data": { "kind": "linspace", "start": 0.05, "end": 0.3 } },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 4, 2],
+          "data": { "kind": "cycle", "values": [78, 211, 39, 185, 124, 104, 154] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "linspace", "start": 0.04, "end": 0.22 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "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": {
+        "inputT": { "dtype": "float32", "shape": [1, 2], "data": { "kind": "values", "values": [1.0, 1.0] } },
+        "routerT": { "dtype": "float32", "shape": [1, 2], "data": { "kind": "values", "values": [0.0, 0.0] } },
+        "fc1T": { "dtype": "uint8", "shape": [2, 2, 1], "data": { "kind": "values", "values": [136, 136, 170, 170] } },
+        "fc1ScalesT": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "constant", "value": 1.0 } },
+        "fc2T": { "dtype": "uint8", "shape": [2, 2, 1], "data": { "kind": "values", "values": [136, 136, 170, 204] } },
+        "fc2ScalesT": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "constant", "value": 1.0 } }
+      },
+      "outputs": {
+        "outputT": {
+          "dtype": "float32",
+          "shape": [1, 2],
+          "data": { "kind": "values", "values": [16.0, 32.0] },
+          "tolerance": 0
+        }
+      }
+    },
+    {
+      "name": "q8_colwise_swiglu",
+      "provenance": {
+        "notes": "Covers 8-bit integer weights and the supported interleaved SwiGLU mode with non-default alpha and beta."
+      },
+      "attrs": {
+        "activation_alpha": 1.1,
+        "activation_beta": 0.05,
+        "activation_type": "swiglu",
+        "expert_weight_bits": 8,
+        "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, 6, 4],
+          "data": { "kind": "cycle", "values": [131, 96, 174, 118, 205, 83, 149] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 6],
+          "data": { "kind": "linspace", "start": 0.006, "end": 0.042 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 4, 3],
+          "data": { "kind": "cycle", "values": [149, 83, 205, 118, 174, 96, 131] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "linspace", "start": 0.008, "end": 0.037 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "q4_blockwise_topk2_swiglu",
+      "provenance": {
+        "notes": "Two complete 16-element quantization blocks on both reductions, with two selected experts and stable top-k softmax mixing."
+      },
+      "attrs": {
+        "activation_type": "swiglu",
+        "block_size": 16,
+        "k": 2,
+        "normalize_routing_weights": 1,
+        "swiglu_fusion": 1,
+        "weights_prepacked": 0
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [4, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.37, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [4, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.23, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [3, 64, 16],
+          "data": { "kind": "cycle", "values": [154, 104, 124, 185, 39, 211, 78] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 64, 2],
+          "data": { "kind": "linspace", "start": 0.04, "end": 0.33 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 32, 16],
+          "data": { "kind": "cycle", "values": [78, 211, 39, 185, 124, 104, 154] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 32, 2],
+          "data": { "kind": "linspace", "start": 0.06, "end": 0.27 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [4, 32], "tolerance": 0.00005 } }
+    },
+    {
+      "name": "q4_rank3_input",
+      "provenance": {
+        "notes": "Locks the standard rank-3 input form and flattening of batch and sequence into router rows."
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 2, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.21, "cosStep": 0.13, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [4, 2],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.17, "scale": 0.8 }
+        },
+        "fc1T": { "dtype": "uint8", "shape": [2, 4, 2], "data": { "kind": "cycle", "values": [25, 118, 203, 71, 154] } },
+        "fc1ScalesT": { "dtype": "float32", "shape": [2, 4], "data": { "kind": "linspace", "start": 0.04, "end": 0.2 } },
+        "fc2T": { "dtype": "uint8", "shape": [2, 4, 2], "data": { "kind": "cycle", "values": [97, 44, 218, 131, 76] } },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "linspace", "start": 0.03, "end": 0.18 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 2, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "q8_swiglu_limit",
+      "provenance": { "notes": "Exercises the optional finite SwiGLU clamp and the explicit 8-bit path." },
+      "attrs": {
+        "activation_alpha": 0.9,
+        "activation_beta": -0.02,
+        "activation_type": "swiglu",
+        "expert_weight_bits": 8,
+        "swiglu_fusion": 1,
+        "swiglu_limit": 0.75
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [2.0, -3.0, 1.0, 4.0, -1.0, 0.5, 3.0, -2.0] }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "data": { "kind": "values", "values": [1.0, 0.0, -1.0, 2.0] }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [2, 8, 4],
+          "data": { "kind": "cycle", "values": [131, 96, 174, 118, 205, 83, 149] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [2, 8],
+          "data": { "kind": "linspace", "start": 0.02, "end": 0.12 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [2, 4, 4],
+          "data": { "kind": "cycle", "values": [149, 83, 205, 118, 174, 96, 131] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "linspace", "start": 0.03, "end": 0.11 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "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 }
+      }
+    },
+    {
+      "name": "empty_zero_dim",
+      "provenance": {
+        "notes": "Zero tokens in the rank-2 form must retain the routed quantized shader identities while every output dispatch is elided."
+      },
+      "inputs": {
+        "inputT": { "dtype": "float32", "shape": [0, 4], "data": { "kind": "values", "values": [] } },
+        "routerT": { "dtype": "float32", "shape": [0, 3], "data": { "kind": "values", "values": [] } },
+        "fc1T": { "dtype": "uint8", "shape": [3, 4, 2], "data": { "kind": "constant", "value": 120 } },
+        "fc1ScalesT": { "dtype": "float32", "shape": [3, 4], "data": { "kind": "constant", "value": 0.05 } },
+        "fc2T": { "dtype": "uint8", "shape": [3, 4, 2], "data": { "kind": "constant", "value": 130 } },
+        "fc2ScalesT": { "dtype": "float32", "shape": [3, 4], "data": { "kind": "constant", "value": 0.05 } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [0, 4], "data": { "kind": "values", "values": [] }, "tolerance": 0 }
+      }
+    },
+    {
+      "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": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.17, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [2, 128, 64],
+          "data": { "kind": "cycle", "values": [154, 104, 124, 185, 39, 211, 78] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "linspace", "start": 0.02, "end": 0.19 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [2, 128, 64],
+          "data": { "kind": "cycle", "values": [78, 211, 39, 185, 124, 104, 154] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "linspace", "start": 0.03, "end": 0.21 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00005 } }
+    },
+    {
+      "name": "q4_blockwise_swiglu_topk2_deep_reduction",
+      "provenance": {
+        "notes": "Interleaved SwiGLU with two selected experts over a 128-long reduction and four 32-wide quantization blocks per row."
+      },
+      "attrs": {
+        "activation_type": "swiglu",
+        "block_size": 32,
+        "k": 2,
+        "normalize_routing_weights": 1,
+        "swiglu_fusion": 1,
+        "weights_prepacked": 0
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.37, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.23, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [3, 256, 64],
+          "data": { "kind": "cycle", "values": [97, 182, 45, 213, 131, 68, 159, 22, 241] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 256, 4],
+          "data": { "kind": "linspace", "start": 0.04, "end": 0.33 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 128, 64],
+          "data": { "kind": "cycle", "values": [241, 22, 159, 68, 131, 213, 45, 182, 97] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 128, 4],
+          "data": { "kind": "linspace", "start": 0.06, "end": 0.27 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 128], "tolerance": 0.00005 } }
+    },
+    {
+      "name": "q8_swiglu_deep_reduction",
+      "provenance": {
+        "notes": "Eight-bit codes over the same 128-long reduction: one stored word per logical column instead of a shared nibble pair."
+      },
+      "attrs": { "activation_type": "swiglu", "expert_weight_bits": 8, "swiglu_fusion": 1, "weights_prepacked": 0 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.11, "scale": 0.35 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.47, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [2, 256, 128],
+          "data": { "kind": "cycle", "values": [113, 26, 201, 88, 167, 54, 229, 141, 73] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [2, 256],
+          "data": { "kind": "linspace", "start": 0.01, "end": 0.09 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [2, 128, 128],
+          "data": { "kind": "cycle", "values": [73, 141, 229, 54, 167, 88, 201, 26, 113] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "linspace", "start": 0.02, "end": 0.08 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00005 } }
+    },
+    {
+      "name": "q8_swiglu_decode_gemv_odd_reduction",
+      "provenance": {
+        "notes": "An 8-bit decode on the cooperative GEMV schedule with an ODD hidden size. The reduction steps two columns at a time -- one stored word per column at 8 bits, one word for both at 4 -- so an odd length leaves a final column that only 8-bit weights can produce."
+      },
+      "attrs": { "activation_type": "swiglu", "expert_weight_bits": 8, "swiglu_fusion": 1, "weights_prepacked": 0 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [1, 129],
+          "data": { "kind": "fillFloat32", "sinStep": 0.261, "cosStep": 0.11, "scale": 0.35 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [1, 2],
+          "data": { "kind": "fillFloat32", "sinStep": 0.341, "cosStep": 0.47, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [2, 256, 129],
+          "data": { "kind": "cycle", "values": [113, 26, 201, 88, 167, 54, 229, 141, 73] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [2, 256],
+          "data": { "kind": "linspace", "start": 0.01, "end": 0.09 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [2, 129, 128],
+          "data": { "kind": "cycle", "values": [73, 141, 229, 54, 167, 88, 201, 26, 113] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [2, 129],
+          "data": { "kind": "linspace", "start": 0.02, "end": 0.08 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 129], "tolerance": 0.00005 } }
+    },
+    {
+      "name": "q4_grouped_prefill_relu",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so one dequantized weight tile is reused across a tile of slots."
+      },
+      "attrs": { "weights_prepacked": 0 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.33999999999999997, "cosStep": 0.24, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [3, 64, 32],
+          "data": { "kind": "cycle", "values": [97, 182, 45, 213, 131, 68, 159, 22, 241] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "linspace", "start": 0.02, "end": 0.21 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 64, 32],
+          "data": { "kind": "cycle", "values": [241, 22, 159, 68, 131, 213, 45, 182, 97] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "linspace", "start": 0.03, "end": 0.18 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 64], "tolerance": 0.00005 } }
+    },
+    {
+      "name": "q4_grouped_prefill_swiglu_topk2",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so one dequantized weight tile is reused across a tile of slots."
+      },
+      "attrs": {
+        "activation_type": "swiglu",
+        "k": 2,
+        "normalize_routing_weights": 1,
+        "swiglu_fusion": 1,
+        "weights_prepacked": 0
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [48, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.37, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [48, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.4, "cosStep": 0.32, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [3, 128, 32],
+          "data": { "kind": "cycle", "values": [97, 182, 45, 213, 131, 68, 159, 22, 241] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "linspace", "start": 0.02, "end": 0.21 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 64, 32],
+          "data": { "kind": "cycle", "values": [241, 22, 159, 68, 131, 213, 45, 182, 97] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "linspace", "start": 0.03, "end": 0.18 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [48, 64], "tolerance": 0.00005 } }
+    },
+    {
+      "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",
+        "k": 2,
+        "normalize_routing_weights": 1,
+        "swiglu_fusion": 1,
+        "weights_prepacked": 0
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [32, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.34, "cosStep": 0.28, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [3, 128, 32],
+          "data": { "kind": "cycle", "values": [97, 182, 45, 213, 131, 68, 159, 22, 241] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "linspace", "start": 0.02, "end": 0.21 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 64, 32],
+          "data": { "kind": "cycle", "values": [241, 22, 159, 68, 131, 213, 45, 182, 97] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "linspace", "start": 0.03, "end": 0.18 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [32, 64], "tolerance": 0.00005 } }
+    },
+    {
+      "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",
+        "block_size": 32,
+        "k": 2,
+        "normalize_routing_weights": 1,
+        "swiglu_fusion": 1,
+        "weights_prepacked": 0,
+        "expert_weight_bits": 8
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [48, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [48, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.44, "cosStep": 0.12000000000000001, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "uint8",
+          "shape": [3, 128, 64],
+          "data": { "kind": "cycle", "values": [97, 182, 45, 213, 131, 68, 159, 22, 241] }
+        },
+        "fc1ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 128, 2],
+          "data": { "kind": "linspace", "start": 0.02, "end": 0.21 }
+        },
+        "fc2T": {
+          "dtype": "uint8",
+          "shape": [3, 64, 64],
+          "data": { "kind": "cycle", "values": [241, 22, 159, 68, 131, 213, 45, 182, 97] }
+        },
+        "fc2ScalesT": {
+          "dtype": "float32",
+          "shape": [3, 64, 2],
+          "data": { "kind": "linspace", "start": 0.03, "end": 0.18 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [48, 64], "tolerance": 0.0003 } }
+    }
+  ]
+}