sync 2e7068faf55e

Browse files

Files changed (14) hide show

README.md +91 -0
build/webgpu/bench.json +473 -0
build/webgpu/expert-group-slots.wgsl.jinja +70 -0
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +26 -0
build/webgpu/moe-ffn-gemv.wgsl.jinja +160 -0
build/webgpu/moe-ffn-grouped.wgsl.jinja +254 -0
build/webgpu/moe-ffn-stage.wgsl.jinja +121 -0
build/webgpu/moe-mix-stage.wgsl.jinja +33 -0
build/webgpu/moe-output-gemv.wgsl.jinja +70 -0
build/webgpu/moe-output-grouped.wgsl.jinja +177 -0
build/webgpu/moe-output-stage.wgsl.jinja +44 -0
build/webgpu/moe-route-stage.wgsl.jinja +76 -0
build/webgpu/test.json +2001 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.MoE
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Mixture of Experts: applies softmax to `router_probs`, routes each token to the top-`k` experts, applies FC1 and `activation_type`, projects through FC2, then sums the selected outputs using their routing probabilities. SwiGLU takes its operands from a separate FC3 (`swiglu_fusion` 0) or a fused FC1 in interleaved (1) or concatenated (2) order; SiLU may also use FC3 as its multiplicative linear projection. This inference package supports float32 and dense routing (`use_sparse_mixer = 0`); float16, bfloat16, and sparse mixing are not implemented. Quantized weights use `com.microsoft.QMoE`.
+See the [ONNX Runtime `MoE` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.MoE) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `input` | `inputT` | `T` | — | — | Token activations, either 2D `(num_tokens, hidden_size)` or 3D `(batch_size, sequence_length, hidden_size)`. | required |
+| `router_probs` | `routerT` | `T` | `2` | — | 2D router logits of shape `(num_tokens, num_experts)`, where `num_tokens` is the product of every leading dimension of `input`. Despite the historical port name, the operator applies a full softmax before top-k selection. | required |
+| `fc1_experts_weights` | `fc1T` | `T` | `3` | — | 3D first-layer expert weights of shape `(num_experts, fusion_size * inter_size, hidden_size)`, where `fusion_size` is 2 for fused SwiGLU (`swiglu_fusion` 1 or 2) and 1 otherwise. | required |
+| `fc1_experts_bias` | `fc1BiasT` | `T` | `2` | — | Optional 2D FC1 bias of shape `(num_experts, fusion_size * inter_size)`. | optional |
+| `fc2_experts_weights` | `fc2T` | `T` | `3` | — | 3D second-layer expert weights of shape `(num_experts, hidden_size, inter_size)`. | required |
+| `fc2_experts_bias` | `fc2BiasT` | `T` | `2` | — | Optional 2D FC2 bias of shape `(num_experts, hidden_size)`, added per expert before that expert's routing weight is applied. | optional |
+| `fc3_experts_weights` | `fc3T` | `T` | `3` | — | Optional 3D third-layer expert weights of shape `(num_experts, inter_size, hidden_size)`. It supplies the separate linear operand for SwiGLU when `swiglu_fusion` is 0, or the multiplicative linear projection for SiLU gating. Other activations do not consume FC3. | optional |
+| `fc3_experts_bias` | `fc3BiasT` | `T` | `2` | — | Optional 2D FC3 bias of shape `(num_experts, inter_size)`. | optional |
+## 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 parameter used by the activation; the schema default is 1. |
+| `activation_beta` | `0` | Beta parameter used by the activation; the schema default is 0. |
+| `activation_type` | `"relu"` | Activation applied to the FC1 projection: `relu`, `gelu`, `silu`, `swiglu`, or `identity`. The schema default is `relu`. |
+| `k` | `1` | Number of experts selected per token; the schema default is 1. |
+| `normalize_routing_weights` | `0` | Whether to normalize the selected routing weights; the schema default is 0. |
+| `swiglu_fusion` | `0` | 0 keeps the SwiGLU operands in separate FC1/FC3 GEMMs, 1 interleaves them in one FC1 row, and 2 concatenates them. The schema default is 0. |
+| `use_sparse_mixer` | `0` | Whether to use sparse-mixer routing. The standard default and only supported value is 0. |
+| `swiglu_limit` | — | Optional SwiGLU clamp limit; omission means no clamp. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32` |
+## 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)
+- [`moe-ffn-gemv.wgsl.jinja`](build/webgpu/moe-ffn-gemv.wgsl.jinja)
+- [`moe-ffn-grouped.wgsl.jinja`](build/webgpu/moe-ffn-grouped.wgsl.jinja)
+- [`moe-ffn-stage.wgsl.jinja`](build/webgpu/moe-ffn-stage.wgsl.jinja)
+- [`moe-mix-stage.wgsl.jinja`](build/webgpu/moe-mix-stage.wgsl.jinja)
+- [`moe-output-gemv.wgsl.jinja`](build/webgpu/moe-output-gemv.wgsl.jinja)
+- [`moe-output-grouped.wgsl.jinja`](build/webgpu/moe-output-grouped.wgsl.jinja)
+- [`moe-output-stage.wgsl.jinja`](build/webgpu/moe-output-stage.wgsl.jinja)
+- [`moe-route-stage.wgsl.jinja`](build/webgpu/moe-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.MoE", { version: 1 });
+const { outputT } = await kernel({
+  inputT: { data: inputTData, shape: [1, 1] },
+  routerT: { data: routerTData, shape: [1, 2] },
+  fc1T: { data: fc1TData, shape: [2, 1, 1] },
+  fc2T: { data: fc2TData, shape: [2, 1, 1] },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,473 @@

+{
+  "op": "com.microsoft.MoE",
+  "tunableSpace": {
+    "workgroupSize": [64, 128],
+    "decodeLanes": [32],
+    "groupThreads": [8, 16],
+    "groupRegM": [1, 2, 4],
+    "groupRegN": [1, 2, 4],
+    "groupTileK": [8, 16, 32]
+  },
+  "cases": [
+    {
+      "name": "moe-f32-t64-h256-e8-k2-i256-swiglu",
+      "attrs": {
+        "k": 2,
+        "activation_type": "swiglu",
+        "swiglu_fusion": 2,
+        "normalize_routing_weights": 1,
+        "activation_alpha": 1.702
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [64, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [64, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [8, 512, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [8, 256, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [64, 256] } }
+    },
+    {
+      "name": "moe-f32-swiglu-t1-h1024-e8-k2-i512-decode",
+      "preset": "smoke",
+      "vars": { "tokens": 1, "hidden": 1024, "experts": 8, "inter": 512, "topK": 2 },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [1, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.33, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.15, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1024] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * 4" }] },
+      "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."
+      }
+    },
+    {
+      "name": "moe-f32-relu-t1-h1024-e8-k1-i1024-decode-relu",
+      "preset": "stress",
+      "vars": { "tokens": 1, "hidden": 1024, "experts": 8, "inter": 1024, "topK": 1 },
+      "attrs": { "k": 1, "activation_type": "relu", "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [1, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.21, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1024] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.topK * 2 * args.hidden * args.inter * 4" }] },
+      "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."
+      }
+    },
+    {
+      "name": "moe-f32-swiglu-t8-h1024-e8-k2-i512-decode-batch8",
+      "preset": "stress",
+      "vars": { "tokens": 8, "hidden": 1024, "experts": 8, "inter": 512, "topK": 2 },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [8, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.21, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [8, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.25, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [8, 1024] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * 4" }] },
+      "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."
+      }
+    },
+    {
+      "name": "moe-f32-swiglu-t256-h1024-e8-k2-i512-prefill",
+      "preset": "stress",
+      "vars": { "tokens": 256, "hidden": 1024, "experts": 8, "inter": 512, "topK": 2 },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [256, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [256, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.49, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.33, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [256, 1024] } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "moe-f32-swiglu-t32-h1024-e8-k2-i512-prefill-small",
+      "preset": "stress",
+      "vars": { "tokens": 32, "hidden": 1024, "experts": 8, "inter": 512, "topK": 2 },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [32, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [32, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.49, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.33, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [32, 1024] } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "moe-f32-swiglu-t64-h1024-e8-k2-i512-prefill-small",
+      "preset": "stress",
+      "vars": { "tokens": 64, "hidden": 1024, "experts": 8, "inter": 512, "topK": 2 },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [64, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.27, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [64, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.49, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.33, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [64, 1024] } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "moe-f32-swiglu-t96-h1024-e8-k2-i512-prefill-small",
+      "preset": "stress",
+      "vars": { "tokens": 96, "hidden": 1024, "experts": 8, "inter": 512, "topK": 2 },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.21, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.49, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.33, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 1024] } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "moe-f32-swiglu-t128-h1024-e8-k2-i512-prefill-small",
+      "preset": "stress",
+      "vars": { "tokens": 128, "hidden": 1024, "experts": 8, "inter": 512, "topK": 2 },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [128, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [128, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.49, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.33, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [128, 1024] } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "moe-f32-swiglu-t48-h1024-e8-k2-i1024-prefill-band",
+      "preset": "stress",
+      "vars": { "tokens": 48, "hidden": 1024, "experts": 8, "inter": 1024, "topK": 2 },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [48, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [48, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.47, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [8, 2048, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.33, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [8, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [48, 1024] } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "moe-f32-swiglu-t96-h1024-e16-k2-i512-prefill-band",
+      "preset": "stress",
+      "vars": { "tokens": 96, "hidden": 1024, "experts": 16, "inter": 512, "topK": 2 },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.53, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [16, 1024, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.33, "cosStep": 0.23, "scale": 0.05 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [16, 1024, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.41, "scale": 0.05 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 1024] } },
+      "bench": {
+        "metrics": [{ "type": "gflops", "value": "2 * args.tokens * args.topK * args.hidden * args.inter * 3" }]
+      }
+    },
+    {
+      "name": "moe-f32-qwen3-moe-decode-t1-h2048-e32-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 32 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."
+      },
+      "vars": { "tokens": 1, "hidden": 2048, "experts": 32, "inter": 768, "topK": 8 },
+      "attrs": { "k": 8, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 3011, "scale": 0.5 },
+        "routerT": { "shape": [1, 32], "dtype": "float32", "dist": "normal", "seed": 3012, "scale": 0.7 },
+        "fc1T": { "shape": [32, 1536, 2048], "dtype": "float32", "dist": "normal", "seed": 3013, "scale": 0.05 },
+        "fc2T": { "shape": [32, 2048, 768], "dtype": "float32", "dist": "normal", "seed": 3014, "scale": 0.05 }
+      },
+      "outputs": { "outputT": { "shape": [1, 2048], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * 4" }] }
+    },
+    {
+      "name": "moe-f32-qwen3-moe-prefill-t256-h2048-e32-k8-i768",
+      "preset": "model",
+      "provenance": {
+        "notes": "Qwen3-MoE class defaults at a 256-token prefill chunk; 2048 routed slots put every expert over the grouped schedule's slot floor."
+      },
+      "vars": { "tokens": 256, "hidden": 2048, "experts": 32, "inter": 768, "topK": 8 },
+      "attrs": { "k": 8, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": { "shape": [256, 2048], "dtype": "float32", "dist": "normal", "seed": 3021, "scale": 0.5 },
+        "routerT": { "shape": [256, 32], "dtype": "float32", "dist": "normal", "seed": 3022, "scale": 0.7 },
+        "fc1T": { "shape": [32, 1536, 2048], "dtype": "float32", "dist": "normal", "seed": 3023, "scale": 0.05 },
+        "fc2T": { "shape": [32, 2048, 768], "dtype": "float32", "dist": "normal", "seed": 3024, "scale": 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": "moe-f32-olmoe-decode-t1-h2048-e16-k8-i2048",
+      "preset": "model",
+      "provenance": {
+        "notes": "OLMoE class defaults (hidden_size 2048, intermediate_size 2048, num_experts_per_tok 8); expert count capped at 16 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."
+      },
+      "vars": { "tokens": 1, "hidden": 2048, "experts": 16, "inter": 2048, "topK": 8 },
+      "attrs": { "k": 8, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 3031, "scale": 0.5 },
+        "routerT": { "shape": [1, 16], "dtype": "float32", "dist": "normal", "seed": 3032, "scale": 0.7 },
+        "fc1T": { "shape": [16, 4096, 2048], "dtype": "float32", "dist": "normal", "seed": 3033, "scale": 0.05 },
+        "fc2T": { "shape": [16, 2048, 2048], "dtype": "float32", "dist": "normal", "seed": 3034, "scale": 0.05 }
+      },
+      "outputs": { "outputT": { "shape": [1, 2048], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * 4" }] }
+    },
+    {
+      "name": "moe-f32-olmoe-prefill-t128-h2048-e16-k8-i2048",
+      "preset": "model",
+      "provenance": { "notes": "OLMoE class defaults at a 128-token prefill chunk." },
+      "vars": { "tokens": 128, "hidden": 2048, "experts": 16, "inter": 2048, "topK": 8 },
+      "attrs": { "k": 8, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": { "shape": [128, 2048], "dtype": "float32", "dist": "normal", "seed": 3041, "scale": 0.5 },
+        "routerT": { "shape": [128, 16], "dtype": "float32", "dist": "normal", "seed": 3042, "scale": 0.7 },
+        "fc1T": { "shape": [16, 4096, 2048], "dtype": "float32", "dist": "normal", "seed": 3043, "scale": 0.05 },
+        "fc2T": { "shape": [16, 2048, 2048], "dtype": "float32", "dist": "normal", "seed": 3044, "scale": 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": "moe-f32-qwen2-moe-decode-t1-h2048-e20-k4-i1408",
+      "preset": "model",
+      "provenance": {
+        "notes": "Qwen2-MoE class defaults (hidden_size 2048, moe_intermediate_size 1408, num_experts_per_tok 4); expert count capped at 20 of the config's 60. 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."
+      },
+      "vars": { "tokens": 1, "hidden": 2048, "experts": 20, "inter": 1408, "topK": 4 },
+      "attrs": { "k": 4, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 3051, "scale": 0.5 },
+        "routerT": { "shape": [1, 20], "dtype": "float32", "dist": "normal", "seed": 3052, "scale": 0.7 },
+        "fc1T": { "shape": [20, 2816, 2048], "dtype": "float32", "dist": "normal", "seed": 3053, "scale": 0.05 },
+        "fc2T": { "shape": [20, 2048, 1408], "dtype": "float32", "dist": "normal", "seed": 3054, "scale": 0.05 }
+      },
+      "outputs": { "outputT": { "shape": [1, 2048], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.topK * 3 * args.hidden * args.inter * 4" }] }
+    },
+    {
+      "name": "moe-f32-qwen2-moe-prefill-t256-h2048-e20-k4-i1408",
+      "preset": "model",
+      "provenance": { "notes": "Qwen2-MoE class defaults at a 256-token prefill chunk." },
+      "vars": { "tokens": 256, "hidden": 2048, "experts": 20, "inter": 1408, "topK": 4 },
+      "attrs": { "k": 4, "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": { "shape": [256, 2048], "dtype": "float32", "dist": "normal", "seed": 3061, "scale": 0.5 },
+        "routerT": { "shape": [256, 20], "dtype": "float32", "dist": "normal", "seed": 3062, "scale": 0.7 },
+        "fc1T": { "shape": [20, 2816, 2048], "dtype": "float32", "dist": "normal", "seed": 3063, "scale": 0.05 },
+        "fc2T": { "shape": [20, 2048, 1408], "dtype": "float32", "dist": "normal", "seed": 3064, "scale": 0.05 }
+      },
+      "outputs": { "outputT": { "shape": [256, 2048], "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

The diff for this file is too large to render. See raw diff

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,26 @@

+{
+  "name": "com.microsoft.MoE",
+  "id": "_com_microsoft_moe_webgpu_9c7f8e1",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "Oq7rf+LMUf5McimJEVeHF3PM3CEHu7X/yCrhSJdoZq4=",
+      "expert-group-slots.wgsl.jinja": "Ta+3H2FA1qRRzksdkwgKLZiO+AokmuM+B9JVN8cv1EQ=",
+      "manifest.json": "1u7xRTHJz3qXke+74gJ0e7GZQdJ9SCJw/8gpaOkB9U4=",
+      "moe-ffn-gemv.wgsl.jinja": "loPQLZYp+Lkz/jq0ML5cYbMFP+IWHlpULSiZcaOx33s=",
+      "moe-ffn-grouped.wgsl.jinja": "m5rGC8aq9oW2TsQ2P3B2uV7Hm7u90DFORehlxFnSw5c=",
+      "moe-ffn-stage.wgsl.jinja": "uIz5vVs9A8s5K5B2ZDj4fYDX3BGsWUfR8LxThIR95Tk=",
+      "moe-mix-stage.wgsl.jinja": "+deuFyuBn04UeIX8Hk2/ntCHNfM1vgkNDTcsh3uXVdI=",
+      "moe-output-gemv.wgsl.jinja": "5yj+BoTsRHiHW92Q8fVvtojdkY7r1jHupMX+FPVQMBk=",
+      "moe-output-grouped.wgsl.jinja": "qzP5G/WtwF4zOBrtySCLTCBflHuvKlU8dk4/CRIzxe8=",
+      "moe-output-stage.wgsl.jinja": "wtB+CtWxl8XChH4Y4/8LEBdHOznrdsUlAo65tmGEuWw=",
+      "moe-route-stage.wgsl.jinja": "6uVRcLbpzp0Y1a8+jRTHOfvBhfCKpCoVju0dBDCVlgo=",
+      "test.json": "T5bB5QuC2mzb2c0i+BQMaTbJjyyvFaL3ikEc4mFJBaw="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MoE" }
+}

build/webgpu/moe-ffn-gemv.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,160 @@

+{{ env.wgsl.resourceDeclarations }}
+// FC1 (and, where the schema splits them, FC3) projection plus the activation, on the
+// cooperative schedule. The split stage gives one thread a whole HIDDEN-long reduction, so a
+// decode dispatch offers only tokens * TOP_K * INTER invocations; here a workgroup owns ROWS
+// intermediate lanes of one routed slot and LANES lanes split the reduction instead.
+// Both operands of a two-operand activation are reduced in the same loop, so each input element
+// is loaded once and used by both projections.
+const HIDDEN: u32 = {{ hidden }}u;
+const INTER: u32 = {{ inter }}u;
+const FC1_ROWS: u32 = {{ fc1Rows }}u;
+const TOP_K: u32 = {{ topK }}u;
+const LANES: u32 = {{ decodeLanes }}u;
+const ROWS: u32 = {{ decodeRows }}u;
+{% if activation == "gelu" %}fn tanh_safe(x: f32) -> f32 {
+  if (x > 10.0) { return 1.0; }
+  if (x < -10.0) { return -1.0; }
+  return tanh(x);
+}
+fn gelu_tanh(v: f32) -> f32 {
+  return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
+}{% endif %}
+{% set isSwiglu = activation == "swiglu" %}
+{% set secondFromFc3 = hasFc3 and (not isSwiglu or swigluFusion == 0) %}
+{% set hasSecond = isSwiglu or hasFc3 %}
+{% set primaryRow = "i * 2u" if (isSwiglu and swigluFusion == 1) else "i" %}
+fn fc1_row(expert: u32, row: u32) -> u32 {
+  return (expert * FC1_ROWS + row) * HIDDEN;
+}
+{% if activation == "swiglu" %}
+fn swiglu(gate_in: f32, up_in: f32) -> f32 {
+{% if hasSwigluLimit %}
+  // swiglu_limit clamps the gate operand from above and the linear operand to [-limit, limit]
+  // before the product; when the attribute is 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);
+}
+{% 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> primary_partial: array<f32, {{ decodeLanes * decodeRows }}>;
+{% if hasSecond %}
+var<workgroup> secondary_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 expert = route_expert[token * TOP_K + local_slot % TOP_K];
+  let in_range = i < INTER;
+  var primary = 0.0;
+{% if hasSecond %}
+  var secondary = 0.0;
+{% endif %}
+  if (in_range) {
+    let a_base = fc1_row(expert, {{ primaryRow }});
+{% if secondFromFc3 %}
+    let b_base = (expert * INTER + i) * HIDDEN;
+{% elif isSwiglu and swigluFusion == 1 %}
+    // swiglu_fusion 1: one fused FC1 whose row pairs are interleaved on each row of 2 * inter.
+    let b_base = fc1_row(expert, i * 2u + 1u);
+{% elif hasSecond %}
+    // swiglu_fusion 2: one fused FC1 whose two halves are concatenated on each row.
+    let b_base = fc1_row(expert, INTER + i);
+{% endif %}
+    for (var h = lane; h < HIDDEN; h = h + LANES) {
+      let x = input[token * HIDDEN + h];
+      primary = primary + x * fc1_experts_weights[a_base + h];
+{% if secondFromFc3 %}
+      secondary = secondary + x * fc3_experts_weights[b_base + h];
+{% elif hasSecond %}
+      secondary = secondary + x * fc1_experts_weights[b_base + h];
+{% endif %}
+    }
+  }
+  let slot = lid.y * LANES + lane;
+  primary_partial[slot] = primary;
+{% if hasSecond %}
+  secondary_partial[slot] = secondary;
+{% endif %}
+  workgroupBarrier();
+{% if hasSecond %}
+{{ rowlane_fold(["primary_partial", "secondary_partial"]) }}
+{% else %}
+{{ rowlane_fold(["primary_partial"]) }}
+{% endif %}
+  if (lane == 0u && in_range) {
+    let base = lid.y * LANES;
+    // A bias belongs to the whole projection, so it is added once to the folded sum rather
+    // than to any lane's partial.
+{% if hasFc1Bias %}
+    let a = primary_partial[base] + fc1_experts_bias[expert * FC1_ROWS + {{ primaryRow }}];
+{% else %}
+    let a = primary_partial[base];
+{% endif %}
+{% if secondFromFc3 and hasFc3Bias %}
+    let b = secondary_partial[base] + fc3_experts_bias[expert * INTER + i];
+{% elif secondFromFc3 %}
+    let b = secondary_partial[base];
+{% elif hasSecond and hasFc1Bias and swigluFusion == 1 %}
+    let b = secondary_partial[base] + fc1_experts_bias[expert * FC1_ROWS + i * 2u + 1u];
+{% elif hasSecond and hasFc1Bias %}
+    let b = secondary_partial[base] + fc1_experts_bias[expert * FC1_ROWS + INTER + i];
+{% elif hasSecond %}
+    let b = secondary_partial[base];
+{% endif %}
+{% if isSwiglu %}
+    let value = swiglu(a, b);
+{% else %}
+{% if activation == "relu" %}
+    let activated = max(a, 0.0);
+{% elif activation == "gelu" %}
+    // The operator uses the tanh GELU approximation.
+    let activated = gelu_tanh(a);
+{% elif activation == "silu" %}
+    let activated = a / (1.0 + exp(-a));
+{% else %}
+    let activated = a;
+{% endif %}
+{% if hasFc3 %}
+    // SiLU with FC3 is the gated-MLP form: SiLU(FC1) multiplies the linear FC3 projection.
+    let value = activated * b;
+{% else %}
+    let value = activated;
+{% endif %}
+{% endif %}
+    hidden_act[local_slot * INTER + i] = value;
+  }
+}

build/webgpu/moe-ffn-grouped.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,254 @@

+{{ env.wgsl.resourceDeclarations }}
+// FC1 (and, where the schema splits them, FC3) projection plus the activation, on the grouped
+// schedule. The other two schedules read a whole expert weight matrix once per routed slot,
+// which at prefill is the entire cost: for hidden 1024 / inter 512 that is 2 MB of weights per
+// slot, re-read for every slot the expert serves.
+//
+// Here the group stage has already sorted the chunk's slots by expert, so a tile of MTILE rows
+// shares ONE expert -- and therefore one weight tile, staged in workgroup memory and reused by
+// every row.
+const HIDDEN: u32 = {{ hidden }}u;
+const INTER: u32 = {{ inter }}u;
+const FC1_ROWS: u32 = {{ fc1Rows }}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;
+{% if activation == "gelu" %}fn tanh_safe(x: f32) -> f32 {
+  if (x > 10.0) { return 1.0; }
+  if (x < -10.0) { return -1.0; }
+  return tanh(x);
+}
+fn gelu_tanh(v: f32) -> f32 {
+  return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
+}{% endif %}
+{% set isSwiglu = activation == "swiglu" %}
+{% set secondFromFc3 = hasFc3 and (not isSwiglu or swigluFusion == 0) %}
+{% set hasSecond = isSwiglu or hasFc3 %}
+{% set primaryRow = "(col * 2u)" if (isSwiglu and swigluFusion == 1) else "col" %}
+{% if secondFromFc3 %}
+{% set secondRowBase = "(expert * INTER + col)" %}
+{% set secondBuffer = "fc3_experts_weights" %}
+{% elif isSwiglu and swigluFusion == 1 %}
+{% set secondRowBase = "(expert * FC1_ROWS + col * 2u + 1u)" %}
+{% set secondBuffer = "fc1_experts_weights" %}
+{% elif hasSecond %}
+{% set secondRowBase = "(expert * FC1_ROWS + INTER + col)" %}
+{% set secondBuffer = "fc1_experts_weights" %}
+{% endif %}
+{% if activation == "swiglu" %}
+fn swiglu(gate_in: f32, up_in: f32) -> f32 {
+{% if hasSwigluLimit %}
+  // swiglu_limit clamps the gate operand from above and the linear operand to [-limit, limit]
+  // before the product; when the attribute is 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);
+}
+{% 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 hasSecond %}
+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();
+{% macro at(buf, idx) %}{% if groupVec4 %}{{ buf }}[({{ idx }}) / 4u][({{ idx }}) % 4u]{% else %}{{ buf }}[{{ idx }}]{% endif %}{% endmacro %}
+{% macro at4(buf, idx) %}{{ buf }}[({{ idx }}) / 4u]{% endmacro %}
+{% set aPre = "(params.tokenOffset + row_slot[m] / TOP_K) * HIDDEN + k_base + " %}
+{% set bPre = "(expert * FC1_ROWS + " ~ primaryRow ~ ") * HIDDEN + k_base + " %}
+{% set cPre = (secondRowBase ~ " * HIDDEN + k_base + ") if hasSecond else "" %}
+{{ group_tile_loop(at("input", aPre ~ "k"),
+                   at("fc1_experts_weights", bPre ~ "k"),
+                   at(secondBuffer, cPre ~ "k") if hasSecond else "",
+                   "HIDDEN", "INTER", regM, regN,
+                   at4("fc1_experts_weights", bPre ~ "kb") if groupVec4 else "",
+                   (at4(secondBuffer, cPre ~ "kb") if hasSecond else "") if groupVec4 else "",
+                   at4("input", aPre ~ "kb") if groupVec4 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 hasFc1Bias %}
+      let a = acc_{{ r }}_{{ c }} + fc1_experts_bias[expert * FC1_ROWS + {{ primaryRow }}];
+{% else %}
+      let a = acc_{{ r }}_{{ c }};
+{% endif %}
+{% if secondFromFc3 and hasFc3Bias %}
+      let b = acc2_{{ r }}_{{ c }} + fc3_experts_bias[expert * INTER + col];
+{% elif secondFromFc3 %}
+      let b = acc2_{{ r }}_{{ c }};
+{% elif hasSecond and hasFc1Bias and swigluFusion == 1 %}
+      let b = acc2_{{ r }}_{{ c }} + fc1_experts_bias[expert * FC1_ROWS + col * 2u + 1u];
+{% elif hasSecond and hasFc1Bias %}
+      let b = acc2_{{ r }}_{{ c }} + fc1_experts_bias[expert * FC1_ROWS + INTER + col];
+{% elif hasSecond %}
+      let b = acc2_{{ r }}_{{ c }};
+{% endif %}
+{% if isSwiglu %}
+      let value = swiglu(a, b);
+{% else %}
+{% if activation == "relu" %}
+      let activated = max(a, 0.0);
+{% elif activation == "gelu" %}
+      // This operator's GELU mode uses the tanh approximation.
+      let activated = gelu_tanh(a);
+{% elif activation == "silu" %}
+      let activated = a / (1.0 + exp(-a));
+{% else %}
+      let activated = a;
+{% endif %}
+{% if hasFc3 %}
+      // SiLU with FC3 is the gated-MLP form: SiLU(FC1) multiplies the linear FC3 projection.
+      let value = activated * b;
+{% else %}
+      let value = activated;
+{% endif %}
+{% endif %}
+      hidden_act[row_slot[m] * INTER + col] = value;
+    }
+  }
+{% endfor %}
+{% endfor %}
+}

build/webgpu/moe-ffn-stage.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,121 @@

+{{ env.wgsl.resourceDeclarations }}
+// FC1 (and, where the schema splits them, FC3) projection plus the activation. One thread owns
+// one intermediate lane of one routed slot, so the grid is (token, slot, inter) flattened and
+// hidden_act is addressed by the same flat index.
+const HIDDEN: u32 = {{ hidden }}u;
+const INTER: u32 = {{ inter }}u;
+// WGSL rejects a const-expression `/ 0u` outright; the contract keeps INTER positive, and this
+// guard keeps the divisor well-formed for any future degenerate arm.
+const INTER_DIV: u32 = max(1u, INTER);
+const FC1_ROWS: u32 = {{ fc1Rows }}u;
+const TOP_K: u32 = {{ topK }}u;
+const WG: u32 = {{ workgroupSize }}u;
+{% if activation == "gelu" %}fn tanh_safe(x: f32) -> f32 {
+  if (x > 10.0) { return 1.0; }
+  if (x < -10.0) { return -1.0; }
+  return tanh(x);
+}
+fn gelu_tanh(v: f32) -> f32 {
+  return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
+}{% endif %}
+fn fc1_row(expert: u32, row: u32) -> u32 {
+  return (expert * FC1_ROWS + row) * HIDDEN;
+}
+fn project_fc1(expert: u32, row: u32, token: u32) -> f32 {
+  let base = fc1_row(expert, row);
+  var acc = 0.0;
+  for (var h: u32 = 0u; h < HIDDEN; h = h + 1u) {
+    acc = acc + input[token * HIDDEN + h] * fc1_experts_weights[base + h];
+  }
+{% if hasFc1Bias %}
+  return acc + fc1_experts_bias[expert * FC1_ROWS + row];
+{% else %}
+  return acc;
+{% endif %}
+}
+{% if hasFc3 %}
+fn project_fc3(expert: u32, row: u32, token: u32) -> f32 {
+  let base = (expert * INTER + row) * HIDDEN;
+  var acc = 0.0;
+  for (var h: u32 = 0u; h < HIDDEN; h = h + 1u) {
+    acc = acc + input[token * HIDDEN + h] * fc3_experts_weights[base + h];
+  }
+{% if hasFc3Bias %}
+  return acc + fc3_experts_bias[expert * INTER + row];
+{% else %}
+  return acc;
+{% endif %}
+}
+{% endif %}
+{% if activation == "swiglu" %}
+fn swiglu(gate_in: f32, up_in: f32) -> f32 {
+{% if hasSwigluLimit %}
+  // swiglu_limit clamps the gate operand from above and the linear operand to [-limit, limit]
+  // before the product; when the attribute is 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);
+}
+{% 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 * 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];
+{% if activation == "swiglu" %}
+{% if swigluFusion == 0 %}
+  // swiglu_fusion 0: the two operands come from separate GEMMs, FC1 and FC3.
+  let value = swiglu(project_fc1(expert, i, token), project_fc3(expert, i, token));
+{% elif swigluFusion == 1 %}
+  // swiglu_fusion 1: one fused FC1 whose row pairs are interleaved on each row of 2 * inter.
+  let value = swiglu(project_fc1(expert, i * 2u, token), project_fc1(expert, i * 2u + 1u, token));
+{% else %}
+  // swiglu_fusion 2: one fused FC1 whose two halves are concatenated on each row.
+  let value = swiglu(project_fc1(expert, i, token), project_fc1(expert, INTER + i, token));
+{% endif %}
+{% else %}
+  let projected = project_fc1(expert, i, token);
+{% if activation == "relu" %}
+  let activated = max(projected, 0.0);
+{% elif activation == "gelu" %}
+  // The operator uses the tanh GELU approximation.
+  let activated = gelu_tanh(projected);
+{% elif activation == "silu" %}
+  let activated = projected / (1.0 + exp(-projected));
+{% else %}
+  let activated = projected;
+{% endif %}
+{% if hasFc3 %}
+  // SiLU with FC3 is the gated-MLP form: SiLU(FC1) multiplies the linear FC3 projection.
+  let value = activated * project_fc3(expert, i, token);
+{% else %}
+  let value = activated;
+{% endif %}
+{% endif %}
+  hidden_act[(local_token * TOP_K + slot % TOP_K) * INTER + i] = value;
+}

build/webgpu/moe-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/moe-output-gemv.wgsl.jinja ADDED Viewed

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

+{{ env.wgsl.resourceDeclarations }}
+// FC2 projection and the routed sum, on the cooperative schedule. A workgroup owns ROWS output
+// columns of one token and LANES lanes split the intermediate axis.
+const HIDDEN: u32 = {{ hidden }}u;
+const INTER: u32 = {{ inter }}u;
+const TOP_K: u32 = {{ topK }}u;
+const LANES: u32 = {{ decodeLanes }}u;
+const ROWS: u32 = {{ decodeRows }}u;
+{% 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 that
+    // slot's own partial here equals scaling the reduced sum, and 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 weight_base = (expert * HIDDEN + out_col) * INTER;
+      let act_base = (local_token * TOP_K + j) * INTER;
+      var part = 0.0;
+      for (var i = lane; i < INTER; i = i + LANES) {
+        part = part + hidden_act[act_base + i] * fc2_experts_weights[weight_base + i];
+      }
+{% if hasFc2Bias %}
+      // The FC2 bias belongs to the expert, so it is added inside the expert's own contribution
+      // and then scaled by that expert's routing weight -- not once to the routed sum. One lane
+      // carries it so the fold adds it exactly once per slot.
+      if (lane == 0u) {
+        part = part + fc2_experts_bias[expert * HIDDEN + out_col];
+      }
+{% endif %}
+      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/moe-output-grouped.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,177 @@

+{{ env.wgsl.resourceDeclarations }}
+// FC2 projection on the grouped schedule, sharing the tile loop with the grouped FC1 stage.
+// The reduction is the intermediate axis and the output width is the hidden axis, so the same
+// expert-homogeneous tile of slots amortizes one staged FC2 weight tile.
+//
+// The result is written PER SLOT rather than per token: a token's routed slots land in
+// different expert tiles, so nothing here can sum them. The mix stage that follows applies each
+// slot's routing weight and collapses the 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;
+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();
+{% macro at(buf, idx) %}{% if groupVec4 %}{{ buf }}[({{ idx }}) / 4u][({{ idx }}) % 4u]{% else %}{{ buf }}[{{ idx }}]{% endif %}{% endmacro %}
+{% macro at4(buf, idx) %}{{ buf }}[({{ idx }}) / 4u]{% endmacro %}
+{% set aPre = "row_slot[m] * INTER + k_base + " %}
+{% set bPre = "(expert * HIDDEN + col) * INTER + k_base + " %}
+{{ group_tile_loop(at("hidden_act", aPre ~ "k"),
+                   at("fc2_experts_weights", bPre ~ "k"),
+                   "", "INTER", "HIDDEN", regM, regN,
+                   at4("fc2_experts_weights", bPre ~ "kb") if groupVec4 else "",
+                   "",
+                   at4("hidden_act", aPre ~ "kb") if groupVec4 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 < HIDDEN) {
+{% if hasFc2Bias %}
+      // The FC2 bias belongs to the expert, so it joins this slot's own contribution before the
+      // mix stage scales it by that slot's routing weight.
+      slot_out[row_slot[m] * HIDDEN + col] = acc_{{ r }}_{{ c }} + fc2_experts_bias[expert * HIDDEN + col];
+{% else %}
+      slot_out[row_slot[m] * HIDDEN + col] = acc_{{ r }}_{{ c }};
+{% endif %}
+    }
+  }
+{% endfor %}
+{% endfor %}
+}

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

	@@ -0,0 +1,44 @@

+{{ env.wgsl.resourceDeclarations }}
+// FC2 projection and the routed sum. One thread owns one output column of one token and walks
+// every routed slot, so the per-slot expert weight is applied to that slot's own projection.
+const HIDDEN: u32 = {{ hidden }}u;
+const INTER: u32 = {{ inter }}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) {
+    let slot = token * TOP_K + j;
+    let expert = route_expert[slot];
+    let weight_base = (expert * HIDDEN + out_col) * INTER;
+    let act_base = (local_token * TOP_K + j) * INTER;
+    var part = 0.0;
+    for (var i: u32 = 0u; i < INTER; i = i + 1u) {
+      part = part + hidden_act[act_base + i] * fc2_experts_weights[weight_base + i];
+    }
+{% if hasFc2Bias %}
+    // The FC2 bias belongs to the expert, so it is added inside the expert's own contribution
+    // and then scaled by that expert's routing weight -- not once to the routed sum.
+    part = part + fc2_experts_bias[expert * HIDDEN + out_col];
+{% endif %}
+    acc = acc + part * route_mix[slot];
+  }
+  output[token * HIDDEN + out_col] = acc;
+}

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

	@@ -0,0 +1,76 @@

+{{ env.wgsl.resourceDeclarations }}
+// One thread per token applies softmax to the router logits, selects TOP_K
+// experts, and writes their optionally renormalized probabilities. Selection is
+// O(TOP_K * EXPERTS) with no scratch. Equal probabilities choose the higher
+// expert index.
+const TOKENS: u32 = {{ tokens }}u;
+const EXPERTS: u32 = {{ experts }}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>) {
+  // gid.y carries the high bits past the per-dimension dispatch limit.
+  let token = gid.x + gid.y * nwg.x * WG;
+  if (token >= TOKENS) {
+    return;
+  }
+  let router_base = token * EXPERTS;
+  var max_logit = router_probs[router_base];
+  for (var e = 1u; e < EXPERTS; e = e + 1u) {
+    max_logit = max(max_logit, router_probs[router_base + e]);
+  }
+  var sum_exp = 0.0;
+  for (var e = 0u; e < EXPERTS; e = e + 1u) {
+    sum_exp = sum_exp + exp(router_probs[router_base + e] - max_logit);
+  }
+  let inv_sum_exp = 1.0 / sum_exp;
+  var experts_taken: array<u32, TOP_K>;
+  var weights: array<f32, TOP_K>;
+  for (var j = 0u; j < TOP_K; j = j + 1u) {
+    // EXPERTS is the sentinel. TOP_K <= EXPERTS guarantees that every pass
+    // finds an untaken expert.
+    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 (experts_taken[p] == e) {
+          taken = true;
+        }
+      }
+      if (!taken) {
+        let candidate = exp(router_probs[router_base + e] - max_logit) * inv_sum_exp;
+        if (best_expert == EXPERTS || candidate >= best_weight) {
+          best_expert = e;
+          best_weight = candidate;
+        }
+      }
+    }
+    experts_taken[j] = best_expert;
+    weights[j] = best_weight;
+  }
+  // Optionally rescale selected probabilities to sum to one. At K == 1 the
+  // single probability becomes exactly one.
+  if (params.normalizeRoutingWeights != 0u) {
+    var sum = 0.0;
+    for (var j = 0u; j < TOP_K; j = j + 1u) {
+      sum = sum + weights[j];
+    }
+    if (sum != 0.0) {
+      for (var j = 0u; j < TOP_K; j = j + 1u) {
+        weights[j] = weights[j] / sum;
+      }
+    }
+  }
+  for (var j = 0u; j < TOP_K; j = j + 1u) {
+    route_expert[token * TOP_K + j] = experts_taken[j];
+    route_mix[token * TOP_K + j] = weights[j];
+  }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,2001 @@

+{
+  "op": "com.microsoft.MoE",
+  "cases": [
+    {
+      "name": "relu_top1_no_bias",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "default activation_type (relu), one expert per token",
+        "notes": "The schema's default activation and nothing optional: no biases, no FC3. Every other case is a delta from this one."
+      },
+      "attrs": { "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "gelu_top1_fc1_bias",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "activation_type gelu with an FC1 bias",
+        "notes": "The tanh-approximate GELU used by ONNX Runtime's MoE kernel, over an FC1 that carries its own per-expert bias. The bias joins the projection before the activation, which is what separates it from a bias on the output side."
+      },
+      "attrs": { "k": 1, "activation_type": "gelu", "normalize_routing_weights": 0 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.13, "scale": 0.25 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "silu_topk2_fc2_bias",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "activation_type silu at k = 2 with an FC2 bias",
+        "notes": "The FC2 bias belongs to the expert, so at k = 2 each selected expert contributes its own bias scaled by its own routing weight. Adding it once to the routed sum instead would be wrong by the difference of the two experts' biases."
+      },
+      "attrs": { "k": 2, "activation_type": "silu", "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.17, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "silu_gate_fc3_no_bias",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "activation_type silu with FC3 as a multiplicative linear projection",
+        "notes": "The pinned CUDA provider's separate FC3 gated-MLP form computes the FC3 projection times SiLU of the FC1 projection."
+      },
+      "attrs": { "k": 1, "activation_type": "silu", "normalize_routing_weights": 0 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "silu_gate_fc3_every_bias",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "all three optional biases present at once",
+        "notes": "FC1, FC2 and FC3 biases together on the provider-defined SiLU gated path."
+      },
+      "attrs": { "k": 2, "activation_type": "silu", "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.13, "scale": 0.25 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.17, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.15 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.23, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "swiglu_fusion2_top1",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "swiglu_fusion 2 (fused, concatenated halves)",
+        "notes": "One fused FC1 whose row block holds the gate half then the linear half. inter_size is the FC1 row count halved, which is what fusion_size 2 means in the schema's shape formula."
+      },
+      "attrs": {
+        "k": 1,
+        "activation_type": "swiglu",
+        "swiglu_fusion": 2,
+        "normalize_routing_weights": 1,
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.33, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.23, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 6, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "swiglu_fusion1_interleaved_topk2",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "swiglu_fusion 1 (fused, interleaved rows)",
+        "notes": "The same fused FC1 as fusion 2 but with the two operands interleaved: lane i reads rows 2i and 2i+1 rather than i and inter+i. Reading the concatenated layout here still lands on the right gate row for lane 0 and on the wrong linear row for it, and on the wrong row for both operands of every later lane."
+      },
+      "attrs": {
+        "k": 2,
+        "activation_type": "swiglu",
+        "swiglu_fusion": 1,
+        "normalize_routing_weights": 1,
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [4, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.33, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [4, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.23, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 6, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.13, "scale": 0.25 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [4, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "swiglu_fusion0_separate_fc3",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "swiglu_fusion 0 (two unfused GEMMs)",
+        "notes": "With the GEMMs unfused the SwiGLU operands come from separate FC1 and FC3 tensors, so fusion_size is 1 and inter_size is the full FC1 row count. FC3 is required in this form and rejected in the fused ones."
+      },
+      "attrs": {
+        "k": 1,
+        "activation_type": "swiglu",
+        "swiglu_fusion": 0,
+        "normalize_routing_weights": 0,
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "swiglu_fusion2_limit_clamped",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "swiglu_limit clamping both operands",
+        "notes": "swiglu_limit clamps the gate operand from above and the linear operand into [-limit, limit] before the product. The limit is small enough here that it actually binds; when the attribute is absent no clamp is applied at all."
+      },
+      "attrs": {
+        "k": 1,
+        "activation_type": "swiglu",
+        "swiglu_fusion": 2,
+        "normalize_routing_weights": 1,
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_limit": 0.35
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.33, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.23, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 6, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.17, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "rank3_batched_relu",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "3D input (batch, sequence, hidden)",
+        "notes": "num_tokens is the product of every leading dimension, so a [1, 3, 4] input routes exactly like the [3, 4] case while the output keeps the 3D shape."
+      },
+      "attrs": { "k": 1, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [1, 3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "topk_equals_expert_count",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "k equal to the expert count",
+        "notes": "Every expert is selected for every token, which is the upper bound the contract allows and the case where the router's taken-exclusion has to visit every index without repeating one."
+      },
+      "attrs": { "k": 3, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "empty_tokens_zero_dim",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "zero-token input",
+        "notes": "A zero-length token axis: every stage dispatches nothing and the output is empty. The route scratch is still allocated at one element so the binding stays well-formed."
+      },
+      "attrs": { "k": 1, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": { "dtype": "float32", "shape": [0, 4], "data": { "kind": "values", "values": [] } },
+        "routerT": { "dtype": "float32", "shape": [0, 3], "data": { "kind": "values", "values": [] } },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [0, 4], "tolerance": 0 } }
+    },
+    {
+      "name": "gelu_fc1_bias_fc2_bias",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "com.microsoft.MoE optional-input combination",
+        "notes": "Each optional input changes the binding layout, so every combination of FC1 bias, FC3 (with or without its own bias) and FC2 bias is a distinct variant. This case is the one that selects fc1bias / no FC3 / fc2bias; without it that variant would ship unexercised."
+      },
+      "attrs": { "k": 1, "activation_type": "gelu", "normalize_routing_weights": 0 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.21, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.13, "scale": 0.25 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.17, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_silu_gate_fc3_fc2_bias",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "com.microsoft.MoE optional-input combination",
+        "notes": "This optional-input layout uses a 128-wide reduction so the cooperative GEMV schedule exercises SiLU with a separate, unbiased FC3 projection and FC2 bias."
+      },
+      "attrs": { "k": 2, "activation_type": "silu", "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.21, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.23, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.17, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_silu_gate_fc3_bias_only",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "com.microsoft.MoE optional-input combination",
+        "notes": "This optional-input layout supplies enough routed slots for the expert-grouped schedule to exercise SiLU with a separately biased FC3 projection and no FC2 bias."
+      },
+      "attrs": { "k": 1, "activation_type": "silu", "normalize_routing_weights": 0 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.25, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.15 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.23, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "silu_gate_fc3_bias_fc2_bias",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "com.microsoft.MoE optional-input combination",
+        "notes": "Each optional input changes the binding layout, so every combination of FC1 bias, FC3 (with or without its own bias) and FC2 bias is a distinct variant. This case is the one that selects no FC1 bias / FC3 with bias / fc2bias; without it that variant would ship unexercised."
+      },
+      "attrs": { "k": 2, "activation_type": "silu", "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.25, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.27, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.17, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.15 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.23, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "silu_fc1_bias_gate_fc3",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "com.microsoft.MoE optional-input combination",
+        "notes": "Each optional input changes the binding layout, so every combination of FC1 bias, FC3 (with or without its own bias) and FC2 bias is a distinct variant. This case is the one that selects fc1bias / FC3 without bias / no FC2 bias; without it that variant would ship unexercised."
+      },
+      "attrs": { "k": 1, "activation_type": "silu", "normalize_routing_weights": 0 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.27, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.29, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.13, "scale": 0.25 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "silu_fc1_bias_gate_fc3_fc2_bias",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "com.microsoft.MoE optional-input combination",
+        "notes": "Each optional input changes the binding layout, so every combination of FC1 bias, FC3 (with or without its own bias) and FC2 bias is a distinct variant. This case is the one that selects fc1bias / FC3 without bias / fc2bias; without it that variant would ship unexercised."
+      },
+      "attrs": { "k": 2, "activation_type": "silu", "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.31, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.13, "scale": 0.25 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.17, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "ort_softmax_tie_break_top1",
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc",
+        "test": "full-softmax routing and equal-probability pair ordering",
+        "notes": "Equal zero logits become probabilities [0.5, 0.5]. ONNX Runtime selects the higher expert index on the tie; that expert emits 6, and the non-normalized route probability scales the result to exactly 3. This also exercises the schema defaults k=1, activation_type=relu, and normalize_routing_weights=0."
+      },
+      "inputs": {
+        "inputT": { "dtype": "float32", "shape": [1, 1], "data": { "kind": "values", "values": [1.0] } },
+        "routerT": { "dtype": "float32", "shape": [1, 2], "data": { "kind": "values", "values": [0.0, 0.0] } },
+        "fc1T": { "dtype": "float32", "shape": [2, 1, 1], "data": { "kind": "values", "values": [1.0, 2.0] } },
+        "fc2T": { "dtype": "float32", "shape": [2, 1, 1], "data": { "kind": "values", "values": [1.0, 3.0] } }
+      },
+      "outputs": {
+        "outputT": {
+          "dtype": "float32",
+          "shape": [1, 1],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [3.0] }
+        }
+      }
+    },
+    {
+      "name": "silu_fc1_bias_gate_fc3_bias",
+      "provenance": {
+        "source": "onnxruntime/docs/ContribOperators.md#com.microsoft.MoE",
+        "test": "com.microsoft.MoE optional-input combination",
+        "notes": "Each optional input changes the binding layout, so every combination of FC1 bias, FC3 (with or without its own bias) and FC2 bias is a distinct variant. This case is the one that selects fc1bias / FC3 with bias / no FC2 bias; without it that variant would ship unexercised."
+      },
+      "attrs": { "k": 1, "activation_type": "silu", "normalize_routing_weights": 0 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.33, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.13, "scale": 0.25 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.16 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.15 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.23, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1plain_fc3none_fc2plain_identity",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": { "activation_type": "identity", "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.31, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.26, "cosStep": 0.24, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.36, "scale": 0.12 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.09, "cosStep": 0.28, "scale": 0.12 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1plain_fc3none_fc2bias_swiglu2",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": {
+        "k": 2,
+        "activation_type": "swiglu",
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_fusion": 2,
+        "normalize_routing_weights": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.083, "cosStep": 0.299, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.273, "cosStep": 0.229, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.123, "cosStep": 0.349, "scale": 0.12 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.103, "cosStep": 0.269, "scale": 0.12 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.233, "cosStep": 0.319, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1bias_fc3none_fc2plain_swiglu1",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 1, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.096, "cosStep": 0.288, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.286, "cosStep": 0.218, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.136, "cosStep": 0.338, "scale": 0.12 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.186, "cosStep": 0.418, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.116, "cosStep": 0.258, "scale": 0.12 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1bias_fc3none_fc2bias_swiglu2",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.109, "cosStep": 0.277, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.299, "cosStep": 0.207, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.149, "cosStep": 0.327, "scale": 0.12 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.199, "cosStep": 0.407, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.129, "cosStep": 0.247, "scale": 0.12 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.259, "cosStep": 0.297, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1plain_fc3plain_fc2plain_swiglu0",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 0, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.122, "cosStep": 0.266, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.312, "cosStep": 0.196, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.162, "cosStep": 0.316, "scale": 0.12 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.142, "cosStep": 0.236, "scale": 0.12 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.182, "cosStep": 0.356, "scale": 0.12 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1plain_fc3plain_fc2bias_swiglu0",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": {
+        "k": 2,
+        "activation_type": "swiglu",
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_fusion": 0,
+        "normalize_routing_weights": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.135, "cosStep": 0.255, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.325, "cosStep": 0.185, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.175, "cosStep": 0.305, "scale": 0.12 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.155, "cosStep": 0.225, "scale": 0.12 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.285, "cosStep": 0.275, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.195, "cosStep": 0.345, "scale": 0.12 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1plain_fc3biased_fc2plain_swiglu0",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": {
+        "activation_type": "swiglu",
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_fusion": 0,
+        "normalize_routing_weights": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.148, "cosStep": 0.244, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.338, "cosStep": 0.174, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.188, "cosStep": 0.294, "scale": 0.12 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.168, "cosStep": 0.214, "scale": 0.12 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.208, "cosStep": 0.334, "scale": 0.12 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.258, "cosStep": 0.194, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1plain_fc3biased_fc2bias_swiglu0",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 0, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.161, "cosStep": 0.233, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.351, "cosStep": 0.163, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.201, "cosStep": 0.283, "scale": 0.12 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.181, "cosStep": 0.203, "scale": 0.12 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.311, "cosStep": 0.253, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.221, "cosStep": 0.323, "scale": 0.12 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.271, "cosStep": 0.183, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1bias_fc3plain_fc2plain_swiglu0",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": {
+        "k": 2,
+        "activation_type": "swiglu",
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_fusion": 0,
+        "normalize_routing_weights": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.174, "cosStep": 0.222, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.364, "cosStep": 0.152, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.214, "cosStep": 0.272, "scale": 0.12 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.264, "cosStep": 0.352, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.194, "cosStep": 0.192, "scale": 0.12 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.234, "cosStep": 0.312, "scale": 0.12 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1bias_fc3plain_fc2bias_swiglu0",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 0, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.187, "cosStep": 0.211, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.377, "cosStep": 0.141, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.227, "cosStep": 0.261, "scale": 0.12 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.277, "cosStep": 0.341, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.207, "cosStep": 0.181, "scale": 0.12 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.337, "cosStep": 0.231, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.247, "cosStep": 0.301, "scale": 0.12 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1bias_fc3biased_fc2plain_swiglu0",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 0, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.2, "cosStep": 0.2, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.39, "cosStep": 0.13, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.24, "cosStep": 0.25, "scale": 0.12 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.33, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.22, "cosStep": 0.17, "scale": 0.12 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.26, "cosStep": 0.29, "scale": 0.12 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.15, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_fc1bias_fc3biased_fc2bias_swiglu0",
+      "provenance": {
+        "notes": "A 128-long reduction on both projections, deep enough for the cooperative schedule; the 4-wide cases above stay on the one-thread-per-column schedule."
+      },
+      "attrs": {
+        "activation_type": "swiglu",
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_fusion": 0,
+        "normalize_routing_weights": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.213, "cosStep": 0.189, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.403, "cosStep": 0.119, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.253, "cosStep": 0.239, "scale": 0.12 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.303, "cosStep": 0.319, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.233, "cosStep": 0.159, "scale": 0.12 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.363, "cosStep": 0.209, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.273, "cosStep": 0.279, "scale": 0.12 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.323, "cosStep": 0.139, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "identity_activation_top1",
+      "provenance": {
+        "notes": "activation_type identity passes the FC1 projection through unchanged, which is the only declared activation with no other case."
+      },
+      "attrs": { "activation_type": "identity" },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.17, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 5, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.18 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.18 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "deep_reduction_gelu_fc2bias",
+      "provenance": {
+        "notes": "GELU over a 128-long reduction makes the cooperative schedule exercise its tanh-approximation activation path."
+      },
+      "attrs": { "k": 2, "activation_type": "gelu", "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.11, "scale": 0.6 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.27, "cosStep": 0.19, "scale": 0.12 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.25, "cosStep": 0.33, "scale": 0.12 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.38, "cosStep": 0.21, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1plain_fc3none_fc2plain_relu",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": { "k": 2, "activation_type": "relu", "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [48, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.09, "cosStep": 0.29, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [48, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.3, "cosStep": 0.24, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.12, "cosStep": 0.36, "scale": 0.15 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.14, "cosStep": 0.26, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [48, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1plain_fc3none_fc2bias_swiglu2",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": {
+        "activation_type": "swiglu",
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_fusion": 2,
+        "normalize_routing_weights": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.107, "cosStep": 0.281, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.317, "cosStep": 0.231, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.137, "cosStep": 0.351, "scale": 0.15 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.157, "cosStep": 0.251, "scale": 0.15 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.277, "cosStep": 0.321, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1bias_fc3none_fc2plain_swiglu1",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 1, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [48, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.124, "cosStep": 0.272, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [48, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.334, "cosStep": 0.222, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.154, "cosStep": 0.342, "scale": 0.15 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.234, "cosStep": 0.422, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.174, "cosStep": 0.242, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [48, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1bias_fc3none_fc2bias_swiglu2",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 2, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.141, "cosStep": 0.263, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.351, "cosStep": 0.213, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.171, "cosStep": 0.333, "scale": 0.15 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.251, "cosStep": 0.413, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.191, "cosStep": 0.233, "scale": 0.15 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.311, "cosStep": 0.303, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_partial_tiles_gelu",
+      "provenance": {
+        "notes": "Half the routed slots the grouped schedule used to demand, so most of every expert's tile is padding: the tiles are 32 rows and 3 experts share only 64 slots. Covers the band the slot threshold opened, where the store guard drops more rows than it keeps."
+      },
+      "attrs": { "activation_type": "gelu", "normalize_routing_weights": 1, "k": 2 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.173, "cosStep": 0.229, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [32, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.347, "cosStep": 0.163, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.199, "cosStep": 0.271, "scale": 0.15 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.251, "cosStep": 0.413, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.241, "cosStep": 0.179, "scale": 0.15 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.311, "cosStep": 0.303, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [32, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1plain_fc3plain_fc2plain_swiglu0",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 0, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.158, "cosStep": 0.254, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.368, "cosStep": 0.204, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.188, "cosStep": 0.324, "scale": 0.15 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.208, "cosStep": 0.224, "scale": 0.15 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.228, "cosStep": 0.364, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1plain_fc3plain_fc2bias_swiglu0",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": {
+        "k": 2,
+        "activation_type": "swiglu",
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_fusion": 0,
+        "normalize_routing_weights": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [48, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.175, "cosStep": 0.245, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [48, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.385, "cosStep": 0.195, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.205, "cosStep": 0.315, "scale": 0.15 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.225, "cosStep": 0.215, "scale": 0.15 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.345, "cosStep": 0.285, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.245, "cosStep": 0.355, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [48, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1plain_fc3biased_fc2plain_swiglu0",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": {
+        "activation_type": "swiglu",
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_fusion": 0,
+        "normalize_routing_weights": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.192, "cosStep": 0.236, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.402, "cosStep": 0.186, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.222, "cosStep": 0.306, "scale": 0.15 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.242, "cosStep": 0.206, "scale": 0.15 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.262, "cosStep": 0.346, "scale": 0.15 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.322, "cosStep": 0.176, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1plain_fc3biased_fc2bias_swiglu0",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 0, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [48, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.209, "cosStep": 0.227, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [48, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.419, "cosStep": 0.177, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.239, "cosStep": 0.297, "scale": 0.15 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.259, "cosStep": 0.197, "scale": 0.15 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.379, "cosStep": 0.267, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.279, "cosStep": 0.337, "scale": 0.15 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.339, "cosStep": 0.167, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [48, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1bias_fc3plain_fc2plain_swiglu0",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": {
+        "k": 2,
+        "activation_type": "swiglu",
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_fusion": 0,
+        "normalize_routing_weights": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [48, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.226, "cosStep": 0.218, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [48, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.436, "cosStep": 0.168, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.256, "cosStep": 0.288, "scale": 0.15 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.336, "cosStep": 0.368, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.276, "cosStep": 0.188, "scale": 0.15 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.296, "cosStep": 0.328, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [48, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1bias_fc3plain_fc2bias_swiglu0",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": { "activation_type": "swiglu", "swiglu_fusion": 0, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.243, "cosStep": 0.209, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.453, "cosStep": 0.159, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.273, "cosStep": 0.279, "scale": 0.15 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.353, "cosStep": 0.359, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.293, "cosStep": 0.179, "scale": 0.15 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.413, "cosStep": 0.249, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.313, "cosStep": 0.319, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1bias_fc3biased_fc2plain_swiglu0",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": { "k": 2, "activation_type": "swiglu", "swiglu_fusion": 0, "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [48, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.26, "cosStep": 0.2, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [48, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.47, "cosStep": 0.15, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.27, "scale": 0.15 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.35, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.17, "scale": 0.15 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.33, "cosStep": 0.31, "scale": 0.15 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.39, "cosStep": 0.14, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [48, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_fc1bias_fc3biased_fc2bias_swiglu0",
+      "provenance": {
+        "notes": "Enough routed slots per expert to fill grouped tiles, so the expert-grouped schedule reuses one staged weight tile across a tile of slots."
+      },
+      "attrs": {
+        "activation_type": "swiglu",
+        "activation_alpha": 1.702,
+        "activation_beta": 0.05,
+        "swiglu_fusion": 0,
+        "normalize_routing_weights": 1
+      },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.277, "cosStep": 0.191, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.487, "cosStep": 0.141, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.307, "cosStep": 0.261, "scale": 0.15 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.387, "cosStep": 0.341, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.327, "cosStep": 0.161, "scale": 0.15 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.447, "cosStep": 0.231, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.347, "cosStep": 0.301, "scale": 0.15 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.407, "cosStep": 0.131, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_identity_no_fc3",
+      "provenance": {
+        "notes": "The smallest default-tile prefill that meets the grouped schedule's routed-slot threshold, and exercises identity activation without FC3 after expert grouping."
+      },
+      "attrs": { "k": 2, "activation_type": "identity", "normalize_routing_weights": 1 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [48, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.181, "cosStep": 0.217, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [48, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.391, "cosStep": 0.127, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.211, "cosStep": 0.287, "scale": 0.15 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.231, "cosStep": 0.187, "scale": 0.15 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [48, 16], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "grouped_prefill_unaligned_hidden_inter",
+      "provenance": {
+        "notes": "Hidden 34 and intermediate 30 are not divisible by four, so grouped staging uses scalar activation and weight reads. The gated FC3 stream adds a second scalar weight stream, and every K tile is partial against the 16-wide tile."
+      },
+      "attrs": { "activation_type": "silu", "normalize_routing_weights": 1, "k": 2 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [96, 34],
+          "data": { "kind": "fillFloat32", "sinStep": 0.163, "cosStep": 0.239, "scale": 0.4 }
+        },
+        "routerT": {
+          "dtype": "float32",
+          "shape": [96, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.337, "cosStep": 0.173, "scale": 0.7 }
+        },
+        "fc1T": {
+          "dtype": "float32",
+          "shape": [3, 30, 34],
+          "data": { "kind": "fillFloat32", "sinStep": 0.209, "cosStep": 0.281, "scale": 0.15 }
+        },
+        "fc1BiasT": {
+          "dtype": "float32",
+          "shape": [3, 30],
+          "data": { "kind": "fillFloat32", "sinStep": 0.261, "cosStep": 0.423, "scale": 0.2 }
+        },
+        "fc2T": {
+          "dtype": "float32",
+          "shape": [3, 34, 30],
+          "data": { "kind": "fillFloat32", "sinStep": 0.251, "cosStep": 0.189, "scale": 0.15 }
+        },
+        "fc2BiasT": {
+          "dtype": "float32",
+          "shape": [3, 34],
+          "data": { "kind": "fillFloat32", "sinStep": 0.321, "cosStep": 0.313, "scale": 0.2 }
+        },
+        "fc3T": {
+          "dtype": "float32",
+          "shape": [3, 30, 34],
+          "data": { "kind": "fillFloat32", "sinStep": 0.227, "cosStep": 0.197, "scale": 0.15 }
+        },
+        "fc3BiasT": {
+          "dtype": "float32",
+          "shape": [3, 30],
+          "data": { "kind": "fillFloat32", "sinStep": 0.283, "cosStep": 0.359, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [96, 34], "tolerance": 0.00002 } }
+    }
+  ]
+}