sync 2e7068faf55e

Browse files

Files changed (7) hide show

README.md +63 -0
build/webgpu/bench.json +52 -0
build/webgpu/gemm-fast-gelu.wgsl.jinja +186 -0
build/webgpu/gemm-subgroup-matrix.wgsl.jinja +241 -0
build/webgpu/manifest.json +246 -0
build/webgpu/metadata.json +19 -0
build/webgpu/test.json +289 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.GemmFastGelu
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Fuses MatMul, an optional bias, and FastGelu: `Y = FastGelu(X @ W + bias)`. `X` has rank at least 2 with shape `(..., K)`, `W` has shape `(K, N)`, and `bias` has shape `(N)`. The activation runs in the float32 accumulator before the output is narrowed, avoiding an intermediate `(..., N)` tensor. Bfloat16 is not implemented.
+See the [ONNX Runtime `GemmFastGelu` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.GemmFastGelu) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `X` | `X` | `T` | — | — | Left operand of rank 2 or greater with shape `(..., K)`; every leading-axis coordinate identifies a row of the product. | required |
+| `W` | `W` | `T` | `2` | — | Right operand with shape `(K, N)`. | required |
+| `bias` | `bias` | `T` | `1` | — | Optional bias with shape `(N)`, added before the activation. | optional |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `Y` | `Y` | `T` | same as `X` | ONNX MatMul result of `X` and `W` | `FastGelu(X @ W + bias)`, with the same rank and leading dimensions as `X` and a trailing `N`. | required |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+## Device requirements
+Some implementation variants require `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`gemm-fast-gelu.wgsl.jinja`](build/webgpu/gemm-fast-gelu.wgsl.jinja)
+- [`gemm-subgroup-matrix.wgsl.jinja`](build/webgpu/gemm-subgroup-matrix.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.GemmFastGelu", { version: 1 });
+const { Y } = await kernel({ X: { data: XData, shape: [5, 6] }, W: { data: WData, shape: [6, 4] } });
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,52 @@

+{
+  "op": "com.microsoft.GemmFastGelu",
+  "cases": [
+    {
+      "name": "gemmfastgelu-bert-base-b8-s384-h768-i3072",
+      "preset": "smoke",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "X": { "shape": [8, 384, 768], "dtype": "float32", "dist": "normal", "seed": 5200, "scale": 1 },
+        "W": { "shape": [768, 3072], "dtype": "float32", "dist": "normal", "seed": 5201, "scale": 1 },
+        "bias": { "shape": [3072], "dtype": "float32", "dist": "normal", "seed": 5202, "scale": 1 }
+      },
+      "outputs": { "Y": { "shape": [8, 384, 3072], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 3072 * 768 * 3072" }] }
+    },
+    {
+      "name": "gemmfastgelu-f16-bert-base-b8-s384-h768-i3072",
+      "preset": "smoke",
+      "vars": { "dtype": "float16" },
+      "inputs": {
+        "X": { "shape": [8, 384, 768], "dtype": "float16", "dist": "normal", "seed": 5210, "scale": 1 },
+        "W": { "shape": [768, 3072], "dtype": "float16", "dist": "normal", "seed": 5211, "scale": 1 },
+        "bias": { "shape": [3072], "dtype": "float16", "dist": "normal", "seed": 5212, "scale": 1 }
+      },
+      "outputs": { "Y": { "shape": [8, 384, 3072], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 3072 * 768 * 3072" }] }
+    },
+    {
+      "name": "gemmfastgelu-bert-large-b4-s512-h1024-i4096",
+      "preset": "model",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "X": { "shape": [4, 512, 1024], "dtype": "float32", "dist": "normal", "seed": 5220, "scale": 1 },
+        "W": { "shape": [1024, 4096], "dtype": "float32", "dist": "normal", "seed": 5221, "scale": 1 },
+        "bias": { "shape": [4096], "dtype": "float32", "dist": "normal", "seed": 5222, "scale": 1 }
+      },
+      "outputs": { "Y": { "shape": [4, 512, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 2048 * 1024 * 4096" }] }
+    },
+    {
+      "name": "gemmfastgelu-nobias-b8-s384-h768-i3072",
+      "preset": "model",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "X": { "shape": [8, 384, 768], "dtype": "float32", "dist": "normal", "seed": 5230, "scale": 1 },
+        "W": { "shape": [768, 3072], "dtype": "float32", "dist": "normal", "seed": 5231, "scale": 1 }
+      },
+      "outputs": { "Y": { "shape": [8, 384, 3072], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 3072 * 768 * 3072" }] }
+    }
+  ]
+}

build/webgpu/gemm-fast-gelu.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,186 @@

+{% if usesF16 is defined and usesF16 %}enable f16;
+{% endif %}{% set OUT = "f16" if (usesF16 is defined and usesF16) else "f32" %}
+{{ env.wgsl.resourceDeclarations }}
+// Shape-driven, register-blocked f32 GEMM for all four transpose combinations.
+// Runtime M/N/K share one compiled pipeline. A 4x4 thread micro-tile forms a
+// 64x64 output tile with BK=16; transpose modes change only load strides and
+// keep the stored-contiguous axis coalesced. M/N/K tails are bounds-checked.
+// Both tiles are indexed by their own output axis and group four K values per
+// vector word, so the micro-tile accumulates through dot() and one step reads
+// TM + TN words instead of 4 * (TM + TN) scalars.
+{% set gemmEpi = gemmEpilogue if gemmEpilogue is defined else "none" %}
+{% if gemmEpi == "fastgelu" %}
+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)));
+}
+{% elif gemmEpi == "activation" %}
+{% set actMode = gemmActivation | default("") %}
+{% set actA = gemmActAlpha | default(0.0) %}
+{% set actB = gemmActBeta | default(0.0) %}
+fn fused_act(v: f32) -> f32 {
+{% if actMode == "Relu" %}
+  return max(v, 0.0);
+{% endif %}
+{% if actMode == "LeakyRelu" %}
+  return select(v * f32({{ actA }}), v, v >= 0.0);
+{% endif %}
+{% if actMode == "Sigmoid" %}
+  return 1.0 / (1.0 + exp(-v));
+{% endif %}
+{% if actMode == "Tanh" %}
+  return tanh(v);
+{% endif %}
+{% if actMode == "HardSigmoid" %}
+  return clamp(f32({{ actA }}) * v + f32({{ actB }}), 0.0, 1.0);
+{% endif %}
+}
+{% endif %}
+{% set epiOpen = "gelu_tanh(" if gemmEpi == "fastgelu" else ("fused_act(" if gemmEpi == "activation" else "") %}
+{% set epiClose = ")" if gemmEpi != "none" else "" %}
+const BK: u32 = 16u;
+const BM: u32 = 64u;
+const BN: u32 = 64u;
+const TM: u32 = 4u;   // per-thread micro-tile rows
+const TN: u32 = 4u;   // per-thread micro-tile cols
+// Stage tiles at the operands' storage width. For f16 inputs this is
+// bit-identical to widening on shared load, keeps f32 FMA registers, and halves
+// workgroup storage.
+{% set tileT = OUT %}
+const K_VECS: u32 = BK / 4u;
+var<workgroup> tileA: array<array<vec4<{{ tileT }}>, K_VECS>, BM>;   // [64][4]  logical A'[m][k/4]
+var<workgroup> tileB: array<array<vec4<{{ tileT }}>, K_VECS>, BN>;   // [64][4]  logical B'[n][k/4]
+@compute @workgroup_size(16, 16, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>
+) {
+  let li = lid.y * 16u + lid.x;
+  let M = params.M;
+  let N = params.N;
+  let K = params.K;
+  let n_tiles = (N + BN - 1u) / BN;
+  let m_tiles = (M + BM - 1u) / BM;
+  for (var nt: u32 = wg.x; nt < n_tiles; nt = nt + nwg.x) {
+    let nBase = nt * BN;
+    for (var mt: u32 = wg.y; mt < m_tiles; mt = mt + nwg.y) {
+      let mBase = mt * BM;
+      var acc: array<f32, TM * TN>;
+      for (var i: u32 = 0u; i < TM * TN; i = i + 1u) { acc[i] = 0.0; }
+      let numTiles = (K + BK - 1u) / BK;
+      for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+        let kBase = kt * BK;
+        // Cooperative load: one vector word per lane per pass. The lane grid is
+        // laid out so the axis the operand stores contiguously stays coalesced.
+        for (var idx: u32 = li; idx < BM * K_VECS; idx = idx + 256u) {
+{% if transA %}
+          // A stored [K, M]: logical A'[m,k] = a[k*M + m]. Coalesce on m (contiguous).
+          let ar = idx % BM;
+          let ac4 = idx / BM;
+{% else %}
+          // A stored [M, K]: logical A[m,k] = a[m*K + k]. Coalesce on k (contiguous).
+          let ar = idx / K_VECS;
+          let ac4 = idx % K_VECS;
+{% endif %}
+          let am = mBase + ar;
+          let ak = kBase + ac4 * 4u;
+          var aWord = vec4<{{ tileT }}>({{ tileT }}(0.0));
+          if (am < M) {
+{% for component in range(4) %}
+            if (ak + {{ component }}u < K) {
+{% if transA %}
+              aWord[{{ component }}u] = {{ tileT }}(a[(ak + {{ component }}u) * M + am]);
+{% else %}
+              aWord[{{ component }}u] = {{ tileT }}(a[am * K + ak + {{ component }}u]);
+{% endif %}
+            }
+{% endfor %}
+          }
+          tileA[ar][ac4] = aWord;
+        }
+        for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + 256u) {
+{% if transB is defined and transB %}
+          // B stored [N, K]: logical B[k,n] = b[n*K + k]. Coalesce on k (contiguous).
+          let bc = idx / K_VECS;
+          let br4 = idx % K_VECS;
+{% else %}
+          // B stored [K, N]: logical B[k,n] = b[k*N + n]. Coalesce on n (contiguous).
+          let bc = idx % BN;
+          let br4 = idx / BN;
+{% endif %}
+          let bn = nBase + bc;
+          let bk = kBase + br4 * 4u;
+          var bWord = vec4<{{ tileT }}>({{ tileT }}(0.0));
+          if (bn < N) {
+{% for component in range(4) %}
+            if (bk + {{ component }}u < K) {
+{% if transB is defined and transB %}
+              bWord[{{ component }}u] = {{ tileT }}(b[bn * K + bk + {{ component }}u]);
+{% else %}
+              bWord[{{ component }}u] = {{ tileT }}(b[(bk + {{ component }}u) * N + bn]);
+{% endif %}
+            }
+{% endfor %}
+          }
+          tileB[bc][br4] = bWord;
+        }
+        workgroupBarrier();
+{% set regIndent = "        " %}{% set regT = "f32" %}{% set regCast = 1 if tileT != "f32" else 0 %}{{ regIndent }}let aRow = lid.y * TM;
+{{ regIndent }}let bCol = lid.x * TN;
+{{ regIndent }}for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
+{{ regIndent }}  var av: array<vec4<{{ regT }}>, TM>;
+{{ regIndent }}  var bv: array<vec4<{{ regT }}>, TN>;
+{{ regIndent }}  for (var i: u32 = 0u; i < TM; i = i + 1u) { av[i] = {% if regCast %}vec4<{{ regT }}>(tileA[aRow + i][kv]){% else %}tileA[aRow + i][kv]{% endif %}; }
+{{ regIndent }}  for (var j: u32 = 0u; j < TN; j = j + 1u) { bv[j] = {% if regCast %}vec4<{{ regT }}>(tileB[bCol + j][kv]){% else %}tileB[bCol + j][kv]{% endif %}; }
+{{ regIndent }}  for (var i: u32 = 0u; i < TM; i = i + 1u) {
+{{ regIndent }}    for (var j: u32 = 0u; j < TN; j = j + 1u) {
+{{ regIndent }}      acc[i * TN + j] = acc[i * TN + j] + dot(av[i], bv[j]);
+{{ regIndent }}    }
+{{ regIndent }}  }
+{{ regIndent }}}
+        workgroupBarrier();
+      }
+      let m0 = mBase + lid.y * TM;
+      let n0 = nBase + lid.x * TN;
+      for (var ti: u32 = 0u; ti < TM; ti = ti + 1u) {
+        let m = m0 + ti;
+        if (m >= M) { continue; }
+        for (var tj: u32 = 0u; tj < TN; tj = tj + 1u) {
+          let n = n0 + tj;
+          if (n < N) {
+            // alpha == 0 nulls A*B exactly and avoids Inf*0 becoming NaN.
+            let product = select(acc[ti * TN + tj] * params.alpha, 0.0, params.alpha == 0.0);
+{% if hasBias is defined and hasBias %}
+{% if rowBias is defined and rowBias %}
+            let biasValue = f32(bias[n]);
+{% elif scalarBias is defined and scalarBias %}
+            let biasValue = f32(bias[0u]);
+{% elif columnBias is defined and columnBias %}
+            let biasValue = f32(bias[m]);
+{% else %}
+            let biasValue = f32(bias[m * N + n]);
+{% endif %}
+            let biasTerm = select(biasValue * params.beta, 0.0, params.beta == 0.0);
+            y[m * N + n] = {{ OUT }}({{ epiOpen }}product + biasTerm{{ epiClose }});
+{% else %}
+            y[m * N + n] = {{ OUT }}({{ epiOpen }}product{{ epiClose }});
+{% endif %}
+          }
+        }
+      }
+    }
+  }
+}

build/webgpu/gemm-subgroup-matrix.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,241 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+enable subgroups;
+{% if pinSubgroupSize32 %}
+enable subgroup_size_control;
+{% endif %}
+enable chromium_experimental_subgroup_matrix;
+diagnostic(off, chromium.subgroup_matrix_uniformity);
+{{ env.wgsl.resourceDeclarations }}
+{% set operandScalar = mmaScalar %}
+{% set accScalar = "f32" %}
+{% set tileMValue = tileM if tileM is defined else 32 %}
+{% set tileNValue = tileN if tileN is defined else 64 %}
+{% set tileKValue = tileK if tileK is defined else 32 %}
+{% set gemmEpi = gemmEpilogue if gemmEpilogue is defined else "none" %}
+{% set epiOpen = "gelu_tanh(" if gemmEpi == "fastgelu" else ("fused_act(" if gemmEpi == "activation" else "") %}
+{% set epiClose = ")" if gemmEpi != "none" else "" %}
+{% set useDoubleBufferedB = doubleBufferedB is defined and doubleBufferedB %}
+{% set bTransposed = bTransposedStorage if bTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
+{% set columnTileBaseValue = columnTileBase if columnTileBase is defined else 0 %}
+{% set rowTileBaseValue = rowTileBase if rowTileBase is defined else 0 %}
+// Four subgroups cover 32xN as a 2x2 grid, or taller tiles as four row bands.
+{% set subtileCols = (tileNValue / 2)|int %}
+{% set subtileRows = 16 %}
+{% set aTilesPerSubgroup = (subtileRows / 8)|int %}
+{% set bTilesPerSubgroup = (subtileCols / 8)|int %}
+{% set subgroupCount = 4 %}
+{% set scratchBanks = subgroupCount * aTilesPerSubgroup * bTilesPerSubgroup %}
+{% if gemmEpi == "fastgelu" %}
+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)));
+}
+{% elif gemmEpi == "activation" %}
+{% set actMode = gemmActivation | default("") %}
+{% set actA = gemmActAlpha | default(0.0) %}
+{% set actB = gemmActBeta | default(0.0) %}
+fn fused_act(v: f32) -> f32 {
+{% if actMode == "Relu" %}
+  return max(v, 0.0);
+{% endif %}
+{% if actMode == "LeakyRelu" %}
+  return select(v * f32({{ actA }}), v, v >= 0.0);
+{% endif %}
+{% if actMode == "Sigmoid" %}
+  return 1.0 / (1.0 + exp(-v));
+{% endif %}
+{% if actMode == "Tanh" %}
+  return tanh(v);
+{% endif %}
+{% if actMode == "HardSigmoid" %}
+  return clamp(f32({{ actA }}) * v + f32({{ actB }}), 0.0, 1.0);
+{% endif %}
+}
+{% endif %}
+const IN_F:       u32 = {{ inFeatures }}u;
+const OUT_F:      u32 = {{ outFeatures }}u;
+{% set kLoopEndValue = kLoopEnd if kLoopEnd is defined else ("K_LOOP" if tailSafe is defined and tailSafe else "IN_F") %}
+const TILE_COLS:  u32 = {{ tileNValue }}u;
+const TILE_ROWS:  u32 = {{ tileMValue }}u;
+const TILE_K:     u32 = {{ tileKValue }}u;
+const SUB_COLS:   u32 = {{ subtileCols }}u;
+const SUB_ROWS:   u32 = {{ subtileRows }}u;
+const COLUMN_TILE_BASE: u32 = {{ columnTileBaseValue }}u;
+const ROW_TILE_BASE: u32 = {{ rowTileBaseValue }}u;
+var<workgroup> tile_A: array<{{ operandScalar }}, {{ tileMValue }} * {{ tileKValue }}>;
+var<workgroup> tile_B: array<{{ operandScalar }}, {{ (2 if useDoubleBufferedB else 1) * tileNValue }} * {{ tileKValue }}>;
+// Distinct readback banks for every result matrix, so one barrier can publish
+// the full subtile without write-after-read reuse.
+var<workgroup> scratch: array<array<{{ accScalar }}, 64>, {{ scratchBanks }}>;
+fn loadSHMA(tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+  // Clamp the partial-M tail to the last valid row (M-1) instead of zero-filling
+  // it. Some subgroup-matrix implementations let a zero-padded left tile corrupt
+  // the adjacent valid row. Replicating row M-1 avoids that boundary; the
+  // replicated rows are discarded by storeOutput's row_limit guard, so every valid
+  // output row stays bit-identical. M >= 1 is guaranteed by the when-clause args.M > 0.
+  let col: u32 = c_idx * 8u;
+  for (var row_offset: u32 = 0u; row_offset < TILE_ROWS; row_offset += {{ (1024 / tileKValue)|int }}u) {
+  let r: u32 = row + row_offset;
+  let a_global: u32 = min(tile_base + r, params.M - 1u);
+  for (var col_offset: u32 = 0u; col_offset < 8u; col_offset++) {
+    let k: u32 = k_idx + col + col_offset;
+    tile_A[r * TILE_K + col + col_offset] =
+      x[a_global * IN_F + k];
+  }
+  }
+}
+fn loadSHMB(tile_base: u32, k_idx: u32, row: u32, c_idx: u32, bank_offset: u32) {
+{% if bTransposed %}
+  // Generic transB=0 stores B as [K,N]. Stage that native orientation and load
+  // the subgroup right operand without a transpose.
+  let local_idx = row * {{ (tileKValue / 16)|int }}u + c_idx;
+  let vectors_per_k = TILE_COLS / 4u;
+  let total_vectors = TILE_K * vectors_per_k;
+  for (var vector_idx = local_idx; vector_idx < total_vectors; vector_idx += 128u) {
+    let k_local = vector_idx / vectors_per_k;
+    let n4 = (vector_idx - k_local * vectors_per_k) * 4u;
+    for (var component = 0u; component < 4u; component++) {
+      let n_local = n4 + component;
+      let global_n = tile_base + n_local;
+      let global_k = k_idx + k_local;
+      tile_B[bank_offset + k_local * TILE_COLS + n_local] = w[global_k * OUT_F + global_n];
+    }
+  }
+{% else %}
+  let col: u32 = c_idx * 16u;
+  for (var row_offset: u32 = 0u; row_offset < TILE_COLS; row_offset += {{ (2048 / tileKValue)|int }}u) {
+  let b_row: u32 = row + row_offset;
+  // Non-power-of-two tuning widths (48/96) use only a prefix of lanes on the
+  // final row band. Keep those inactive lanes out of tile_B and global memory.
+  if (b_row < TILE_COLS) {
+    let w_global: u32 = tile_base + b_row;
+  for (var i: u32 = 0u; i < 16u; i++) {
+    let k: u32 = k_idx + col + i;
+    tile_B[bank_offset + b_row * TILE_K + col + i] =
+      w[w_global * IN_F + k];
+  }
+  }
+  }
+{% endif %}
+}
+{% set needsColBase = hasBias or (tailSafe is defined and tailSafe) %}
+{% set BIAS_1 = (" + " ~ accScalar ~ "(bias[col_base + col])") if hasBias else "" %}
+{% set BIAS_2 = (" + " ~ accScalar ~ "(bias[col_base + col2])") if hasBias else "" %}
+fn storeOutput(offset: u32{% if needsColBase %}, col_base: u32{% endif %}, row: u32, col: u32, src_slot: u32, row_limit: i32) {
+  if (row_limit > 0 && row < u32(row_limit)) {
+    let col2: u32 = col + 1u;
+    y[offset + row * OUT_F + col]  = {{ yScalar }}({{ epiOpen }}scratch[src_slot][row * 8u + col]{{ BIAS_1 }}{{ epiClose }});
+    y[offset + row * OUT_F + col2] = {{ yScalar }}({{ epiOpen }}scratch[src_slot][row * 8u + col2]{{ BIAS_2 }}{{ epiClose }});
+  }
+}
+@compute @workgroup_size(128, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
+fn main(
+  @builtin(workgroup_id) workgroup_id: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_index) local_idx: u32,
+  @builtin(subgroup_invocation_id) sg_id: u32,
+  @builtin(subgroup_size) sg_size: u32
+) {
+  // 2D-folded M-tile (row) base: workgroup_id.z carries the high bits when the
+  // row tile exceeds the per-dimension dispatch limit.
+  // Reduces to workgroup_id.y when nwg.z == 1; the row_limit guard in storeOutput
+  // (i32(M) - i32(a_global_base + ...)) drops the over-dispatched tail.
+  let a_global_base: u32 = (ROW_TILE_BASE + workgroup_id.y + workgroup_id.z * nwg.y) * TILE_ROWS;
+  let w_global_base: u32 = (COLUMN_TILE_BASE + workgroup_id.x) * TILE_COLS;
+  let subtile_id: u32 = local_idx / sg_size;
+  let subtile_idx: u32 = subtile_id / 2u;
+  let subtile_idy: u32 = subtile_id % 2u;
+  let base_A: u32 = subtile_idy * SUB_ROWS;
+  let base_B: u32 = subtile_idx * SUB_COLS;
+{% for m in range(aTilesPerSubgroup) %}
+{% for n in range(bTilesPerSubgroup) %}
+  var matC{{ m }}{{ n }}: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+{% endfor %}
+{% endfor %}
+  for (var kidx: u32 = 0u; kidx < {{ kLoopEndValue }}; kidx += TILE_K) {
+    loadSHMA(a_global_base, kidx, local_idx / {{ (tileKValue / 8)|int }}u, local_idx % {{ (tileKValue / 8)|int }}u);
+    loadSHMB(w_global_base, kidx, local_idx / {{ (tileKValue / 16)|int }}u, local_idx % {{ (tileKValue / 16)|int }}u, 0u);
+    workgroupBarrier();
+    for (var step: u32 = 0u; step < TILE_K; step += 8u) {
+{% set dynamicATiles = aTilesPerSubgroup if aTilesPerSubgroup is defined else 2 %}
+{% set bTransposed = bTransposedStorage if bTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
+{% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+{% for m in range(dynamicATiles) %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, false, TILE_K);
+{% endfor %}
+{% if bTransposed %}
+      let matrix_b_offset = {{ B_BANK }}step * TILE_COLS + subtile_idx * SUB_COLS;
+{% else %}
+      let matrix_b_offset = {{ B_BANK }}subtile_idx * SUB_COLS * TILE_K + step;
+{% endif %}
+{% for n in range(bTilesPerSubgroup) %}
+{% if bTransposed %}
+      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + {{ n * 8 }}u, false, TILE_COLS);
+{% else %}
+      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + {{ n * 8 }}u * TILE_K, true, TILE_K);
+{% endif %}
+{% endfor %}
+{% for m in range(dynamicATiles) %}
+{% for n in range(bTilesPerSubgroup) %}
+      matC{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, matC{{ m }}{{ n }});
+{% endfor %}
+{% endfor %}
+    }
+    workgroupBarrier();
+  }
+  let matrix_c_offset: u32 = (a_global_base + base_A) * OUT_F + w_global_base + base_B;
+  let row: u32 = sg_id / 4u;
+  let col: u32 = (sg_id % 4u) * 2u;
+  // Stage every result into a distinct bank before one barrier. Tail/bias paths
+  // retain this guarded scalar epilogue: subgroupMatrixStore scatters values
+  // across lanes, so its cross-lane readback must be published before partial-M
+  // guards diverge. The epilogue also handles output conversion and bias.
+  let bank: u32 = subtile_id * {{ aTilesPerSubgroup * bTilesPerSubgroup }}u;
+{% for m in range(aTilesPerSubgroup) %}
+{% for n in range(bTilesPerSubgroup) %}
+  subgroupMatrixStore(&scratch[bank + {{ m * bTilesPerSubgroup + n }}u], 0u, matC{{ m }}{{ n }}, false, 8u);
+{% endfor %}
+{% endfor %}
+  workgroupBarrier();
+{% for m in range(aTilesPerSubgroup) %}
+  let row_limit_{{ m }}: i32 = i32(params.M) - i32(a_global_base + base_A + {{ m * 8 }}u);
+{% for n in range(bTilesPerSubgroup) %}
+  storeOutput(
+    matrix_c_offset + {{ m * 8 }}u * OUT_F + {{ n * 8 }}u{% if needsColBase %},
+    w_global_base + base_B + {{ n * 8 }}u,
+{% else %},
+{% endif %}
+    row,
+    col,
+    bank + {{ m * bTilesPerSubgroup + n }}u,
+    row_limit_{{ m }}
+  );
+{% endfor %}
+{% endfor %}
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,246 @@

+{
+  "domain": "com.microsoft",
+  "name": "GemmFastGelu",
+  "sinceVersion": 1,
+  "description": "Fuses MatMul, an optional bias, and FastGelu: `Y = FastGelu(X @ W + bias)`. `X` has rank at least 2 with shape `(..., K)`, `W` has shape `(K, N)`, and `bias` has shape `(N)`. The activation runs in the float32 accumulator before the output is narrowed, avoiding an intermediate `(..., N)` tensor. Bfloat16 is not implemented.",
+  "inputs": [
+    {
+      "role": "X",
+      "dtype": "T",
+      "description": "Left operand of rank 2 or greater with shape `(..., K)`; every leading-axis coordinate identifies a row of the product."
+    },
+    { "role": "W", "dtype": "T", "rank": 2, "description": "Right operand with shape `(K, N)`." },
+    {
+      "role": "bias",
+      "dtype": "T",
+      "rank": 1,
+      "optional": true,
+      "description": "Optional bias with shape `(N)`, added before the activation."
+    }
+  ],
+  "outputs": [
+    {
+      "role": "Y",
+      "dtype": "T",
+      "rank": "ranks.X",
+      "shape": "matmulShape(shapes.X, shapes.W)",
+      "description": "`FastGelu(X @ W + bias)`, with the same rank and leading dimensions as `X` and a trailing `N`."
+    }
+  ],
+  "typeConstraints": { "T": ["float32", "float16"] },
+  "args": {
+    "X": { "kind": "tensor", "semantic": "X", "role": "input" },
+    "W": { "kind": "tensor", "semantic": "W", "role": "weights" },
+    "bias": { "kind": "tensor", "semantic": "bias", "role": "weights", "required": false },
+    "Y": { "kind": "tensor", "semantic": "Y", "role": "output" }
+  },
+  "derive": {
+    "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
+    "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
+    "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
+    "wave32Effective": "wave32Adapter or pinSubgroupSize32",
+    "kDim": "dim(shapes.W, 0)",
+    "nDim": "dim(shapes.W, 1)",
+    "gfgRows": "numel(shapes.X) / kDim if kDim > 0 else 0",
+    "shapeOk": "ranks.W == 2 and ranks.X >= 2 and kDim > 0 and nDim > 0 and dim(shapes.X, ranks.X - 1) == kDim and ranks.Y == ranks.X and dim(shapes.Y, ranks.Y - 1) == nDim and numel(shapes.Y) == gfgRows * nDim",
+    "dtypeOk": "(tensorDtypes.X == \"float32\" or tensorDtypes.X == \"float16\") and tensorDtypes.W == tensorDtypes.X and tensorDtypes.Y == tensorDtypes.X and f16Ok(dtypes.T)",
+    "biasOk": "ranks.bias == 1 and dim(shapes.bias, 0) == nDim and tensorDtypes.bias == tensorDtypes.X",
+    "baseContract": "shapeOk and dtypeOk",
+    "noBiasContract": "baseContract and not present.bias",
+    "biasContract": "baseContract and present.bias and biasOk",
+    "sgmatOperandBytes": "2 if tensorDtypes.X == \"float16\" and tensorDtypes.W == \"float16\" else 4",
+    "sgmatStorageBytes": "(32 * 32 + 64 * 32) * sgmatOperandBytes + 32 * 64 * 4",
+    "sgmatResourcesFit": "128 <= deviceWorkgroupCap and sgmatStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "sgmatLayoutOk": "gfgRows > 0 and kDim % 32 == 0 and nDim % 64 == 0 and ((tensorDtypes.X == \"float16\" and device.features.has(\"shader-f16\") and gfgRows >= 2) or (tensorDtypes.X == \"float32\" and gfgRows >= 32)) and ceilDiv(gfgRows, 32) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(nDim, 64) <= device.limits.maxComputeWorkgroupsPerDimension",
+    "sgmatContract": "wave32Effective and sgmatLayoutOk and sgmatResourcesFit"
+  },
+  "constants": {
+    "usesF16": "tensorDtypes.X == \"float16\"",
+    "transA": false,
+    "transB": false,
+    "gemmEpilogue": "\"fastgelu\""
+  },
+  "bindingSets": {
+    "noBias": [
+      { "name": "a", "arg": "X", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "b", "arg": "W", "semantic": "W", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "M", "type": "u32", "value": "gfgRows" },
+            { "name": "N", "type": "u32", "value": "nDim" },
+            { "name": "K", "type": "u32", "value": "kDim" },
+            { "name": "alpha", "type": "f32", "value": 1 }
+          ]
+        }
+      }
+    ],
+    "withBias": [
+      { "name": "a", "arg": "X", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "b", "arg": "W", "semantic": "W", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      {
+        "name": "bias",
+        "arg": "bias",
+        "semantic": "bias",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$T"
+      },
+      { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "M", "type": "u32", "value": "gfgRows" },
+            { "name": "N", "type": "u32", "value": "nDim" },
+            { "name": "K", "type": "u32", "value": "kDim" },
+            { "name": "alpha", "type": "f32", "value": 1 },
+            { "name": "beta", "type": "f32", "value": 1 }
+          ]
+        }
+      }
+    ],
+    "sgmatNoBias": [
+      { "name": "x", "arg": "X", "semantic": "a", "buffer": { "type": "read-only-storage" }, "elementType": "$xScalar" },
+      { "name": "w", "arg": "W", "semantic": "b", "buffer": { "type": "read-only-storage" }, "elementType": "$wScalar" },
+      { "name": "y", "arg": "Y", "semantic": "y", "buffer": { "type": "storage" }, "elementType": "$yScalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "M", "type": "u32", "value": "gfgRows" }] }
+      }
+    ],
+    "sgmatBias": [
+      { "name": "x", "arg": "X", "semantic": "a", "buffer": { "type": "read-only-storage" }, "elementType": "$xScalar" },
+      { "name": "w", "arg": "W", "semantic": "b", "buffer": { "type": "read-only-storage" }, "elementType": "$wScalar" },
+      {
+        "name": "bias",
+        "arg": "bias",
+        "semantic": "bias",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$bScalar"
+      },
+      { "name": "y", "arg": "Y", "semantic": "y", "buffer": { "type": "storage" }, "elementType": "$yScalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "M", "type": "u32", "value": "gfgRows" }] }
+      }
+    ]
+  },
+  "variants": [
+    {
+      "id": "sgmat_bias",
+      "priority": 100,
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [
+          { "componentType": "f16", "M": 8, "N": 8, "K": 8 },
+          { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
+        ]
+      },
+      "when": ["biasContract", "sgmatContract"],
+      "constants": {
+        "hasBias": true,
+        "gemmEpilogue": "\"fastgelu\"",
+        "bTransposedStorage": true,
+        "xScalar": "\"f16\" if tensorDtypes.X == \"float16\" else \"f32\"",
+        "wScalar": "\"f16\" if tensorDtypes.W == \"float16\" else \"f32\"",
+        "bScalar": "\"f16\" if tensorDtypes.bias == \"float16\" else \"f32\"",
+        "yScalar": "\"f16\" if tensorDtypes.Y == \"float16\" else \"f32\"",
+        "mmaScalar": "\"f16\" if tensorDtypes.X == \"float16\" and tensorDtypes.W == \"float16\" else \"f32\"",
+        "usesF16": "tensorDtypes.X == \"float16\" or tensorDtypes.W == \"float16\" or tensorDtypes.Y == \"float16\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.SubgroupMatrixBias",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": "sgmatBias",
+          "dispatch": { "workgroups": "ceilDiv(gfgRows, 32)", "x": "ceilDiv(nDim, 64)" }
+        }
+      ]
+    },
+    {
+      "id": "sgmat",
+      "priority": 100,
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [
+          { "componentType": "f16", "M": 8, "N": 8, "K": 8 },
+          { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
+        ]
+      },
+      "when": ["noBiasContract", "sgmatContract"],
+      "constants": {
+        "hasBias": false,
+        "gemmEpilogue": "\"fastgelu\"",
+        "bTransposedStorage": true,
+        "xScalar": "\"f16\" if tensorDtypes.X == \"float16\" else \"f32\"",
+        "wScalar": "\"f16\" if tensorDtypes.W == \"float16\" else \"f32\"",
+        "yScalar": "\"f16\" if tensorDtypes.Y == \"float16\" else \"f32\"",
+        "mmaScalar": "\"f16\" if tensorDtypes.X == \"float16\" and tensorDtypes.W == \"float16\" else \"f32\"",
+        "usesF16": "tensorDtypes.X == \"float16\" or tensorDtypes.W == \"float16\" or tensorDtypes.Y == \"float16\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.SubgroupMatrix",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": "sgmatNoBias",
+          "dispatch": { "workgroups": "ceilDiv(gfgRows, 32)", "x": "ceilDiv(nDim, 64)" }
+        }
+      ]
+    },
+    {
+      "id": "tiled_bias",
+      "priority": 0,
+      "when": ["biasContract"],
+      "constants": { "hasBias": true, "rowBias": true, "scalarBias": false, "columnBias": false },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.Bias",
+          "shader": "gemm-fast-gelu.wgsl.jinja",
+          "bindings": "withBias",
+          "dispatch": {
+            "x": "min(ceilDiv(nDim, 64), device.limits.maxComputeWorkgroupsPerDimension)",
+            "y": "min(ceilDiv(gfgRows, 64), device.limits.maxComputeWorkgroupsPerDimension)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "tiled",
+      "priority": 0,
+      "when": ["noBiasContract"],
+      "constants": { "hasBias": false },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu",
+          "shader": "gemm-fast-gelu.wgsl.jinja",
+          "bindings": "noBias",
+          "dispatch": {
+            "x": "min(ceilDiv(nDim, 64), device.limits.maxComputeWorkgroupsPerDimension)",
+            "y": "min(ceilDiv(gfgRows, 64), device.limits.maxComputeWorkgroupsPerDimension)"
+          }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,19 @@

+{
+  "name": "com.microsoft.GemmFastGelu",
+  "id": "_com_microsoft_gemmfastgelu_webgpu_372327f",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "CPFUVoaXFHE1nKK5oZMKmB0LohpzyltnfuRq/70nJAk=",
+      "gemm-fast-gelu.wgsl.jinja": "IkB53JqVNn9Xezw0YcKYDvm9IwDvoZNnSeA/UY4Eci4=",
+      "gemm-subgroup-matrix.wgsl.jinja": "2PtV5JZcKiR+CKWoCbReeevzatqfe9l/O6U5wwbDgoU=",
+      "manifest.json": "bcpv0SoGGVwU23+3nF7YJ2BSX80r64AInCEZgAuCRoI=",
+      "test.json": "WNyAFlYnnkTg3N4r9ja8Kh2HwlO8jr7EncOBBa9oI48="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GemmFastGelu" }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,289 @@

+{
+  "op": "com.microsoft.GemmFastGelu",
+  "cases": [
+    {
+      "name": "rank2_bias_f32",
+      "provenance": {
+        "notes": "A plain (M, K) x (K, N) with bias, both dimensions past one 64-wide tile so the kernel's tile loop and its bounds checks both run."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [96, 80],
+          "data": { "kind": "fillFloat32", "scale": 0.9, "sinStep": 0.18000000000000002, "cosStep": 0.32 }
+        },
+        "W": {
+          "dtype": "float32",
+          "shape": [80, 72],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.12, "cosStep": 0.24000000000000002 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [72],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.3, "cosStep": 0.14 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [96, 72], "tolerance": 0.00002, "relTolerance": 0.0002 } }
+    },
+    {
+      "name": "rank2_nobias_f32",
+      "provenance": { "notes": "The same shape without a bias, which takes the other binding set." },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [96, 80],
+          "data": { "kind": "fillFloat32", "scale": 0.9, "sinStep": 0.19, "cosStep": 0.33 }
+        },
+        "W": {
+          "dtype": "float32",
+          "shape": [80, 72],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.13, "cosStep": 0.25 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [96, 72], "tolerance": 0.00002, "relTolerance": 0.0002 } }
+    },
+    {
+      "name": "rank3_batched_bias_f32",
+      "provenance": {
+        "notes": "A rank-3 X: every leading axis is another row, so a (2, 48, 64) input is a 96-row product rather than a batched matmul."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [2, 48, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.9, "sinStep": 0.2, "cosStep": 0.33999999999999997 }
+        },
+        "W": {
+          "dtype": "float32",
+          "shape": [64, 40],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.14, "cosStep": 0.26 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [40],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.31999999999999995, "cosStep": 0.16 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 48, 40], "tolerance": 0.00002, "relTolerance": 0.0002 } }
+    },
+    {
+      "name": "rank4_batched_nobias_f32",
+      "provenance": {
+        "notes": "Rank 4, no bias -- the leading-axis collapse must not depend on how many axes there are."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [2, 3, 16, 32],
+          "data": { "kind": "fillFloat32", "scale": 0.9, "sinStep": 0.21000000000000002, "cosStep": 0.35 }
+        },
+        "W": {
+          "dtype": "float32",
+          "shape": [32, 24],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.15, "cosStep": 0.27 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 16, 24], "tolerance": 0.00002, "relTolerance": 0.0002 } }
+    },
+    {
+      "name": "small_unaligned_bias_f32",
+      "provenance": {
+        "notes": "M, K and N all smaller than one tile and none a multiple of the 4x4 micro-tile, so every bounds check is exercised."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [7, 13],
+          "data": { "kind": "fillFloat32", "scale": 0.9, "sinStep": 0.22000000000000003, "cosStep": 0.36 }
+        },
+        "W": {
+          "dtype": "float32",
+          "shape": [13, 11],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.16, "cosStep": 0.28 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [11],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.33999999999999997, "cosStep": 0.18 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [7, 11], "tolerance": 0.00002, "relTolerance": 0.0002 } }
+    },
+    {
+      "name": "single_row_bias_f32",
+      "provenance": { "notes": "One row: the GEMV shape, where the M tile is almost entirely padding." },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [1, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.9, "sinStep": 0.23, "cosStep": 0.37 }
+        },
+        "W": {
+          "dtype": "float32",
+          "shape": [64, 48],
+          "data": {
+            "kind": "fillFloat32",
+            "scale": 0.5,
+            "sinStep": 0.16999999999999998,
+            "cosStep": 0.29000000000000004
+          }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [48],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.35, "cosStep": 0.19 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [1, 48], "tolerance": 0.00002, "relTolerance": 0.0002 } }
+    },
+    {
+      "name": "f16_rank3_bias",
+      "provenance": {
+        "notes": "float16 operands; the product accumulates in f32 and the activation runs before the single narrowing store."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float16",
+          "shape": [2, 32, 48],
+          "data": { "kind": "fillFloat32", "scale": 0.9, "sinStep": 0.24000000000000002, "cosStep": 0.38 }
+        },
+        "W": {
+          "dtype": "float16",
+          "shape": [48, 32],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.18, "cosStep": 0.30000000000000004 }
+        },
+        "bias": {
+          "dtype": "float16",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.36, "cosStep": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 32, 32], "tolerance": 0.006, "relTolerance": 0.006 } }
+    },
+    {
+      "name": "f16_rank2_nobias",
+      "provenance": { "notes": "float16 without a bias." },
+      "inputs": {
+        "X": {
+          "dtype": "float16",
+          "shape": [48, 32],
+          "data": { "kind": "fillFloat32", "scale": 0.9, "sinStep": 0.25, "cosStep": 0.39 }
+        },
+        "W": {
+          "dtype": "float16",
+          "shape": [32, 24],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.19, "cosStep": 0.31 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [48, 24], "tolerance": 0.006, "relTolerance": 0.006 } }
+    },
+    {
+      "name": "pinned_rank2_bias",
+      "provenance": {
+        "notes": "Hand-computed from the schema's FastGelu definition -- 0.5x(1 + tanh(sqrt(2/pi)(x + 0.044715 x^3))) -- applied to a float64 product, independent of the port's tiling and of its shared gelu helper. The operands are large enough that FastGelu is well outside its near-linear region, so the pin is sensitive to the activation itself."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [5, 6],
+          "data": {
+            "kind": "values",
+            "values": [-0.568, 1.402, -1.744, 0.805, -0.684, -0.366, -1.898, -0.274, 0.008, -1.719, 0.344, 1.185, -0.865, 1.863, -0.792, 1.783, 0.703, -0.614, -1.244, 1.186, 1.21, -0.428, -0.236, 1.268, -1.492, -1.469, 0.673, 1.956, -1.883, -0.872]
+          }
+        },
+        "W": {
+          "dtype": "float32",
+          "shape": [6, 4],
+          "data": {
+            "kind": "values",
+            "values": [-0.362, -1.699, 1.664, -0.171, 0.412, -0.373, -1.421, 1.657, 1.194, -0.13, -1.707, -0.827, -0.156, -1.618, 1.339, -0.16, -1.723, 1.951, 0.603, 1.677, 1.217, 0.273, 1.35, 0.713]
+          }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "values", "values": [-0.818, 0.02, 0.585, 0.387] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [5, 4],
+          "data": {
+            "kind": "values",
+            "values": [-0.099208, -0.041387, 0.626313, 2.704125, 0.716692, 7.121855, -0.009083, 1.89771, -0.004677, -0.16984, -0.072533, 4.732758, 3.581808, 2.075684, -0.000021, 2.106721, 1.733585, -0.000052, -0.167767, -0.0]
+          },
+          "tolerance": 0.00002,
+          "relTolerance": 0.0002
+        }
+      }
+    },
+    {
+      "name": "sgmat_aligned_bias_f32",
+      "provenance": {
+        "notes": "M, K and N all meet the subgroup-matrix gate (M >= 32, K % 32 == 0, N % 64 == 0), so this is the fixture that selects sgmat_bias rather than the tiled kernel. Two row tiles and two column tiles past the 32x64 tile, so the tile loop runs and the fused gelu is checked on every store site."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [64, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.7, "sinStep": 0.13, "cosStep": 0.29 }
+        },
+        "W": {
+          "dtype": "float32",
+          "shape": [64, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.21, "cosStep": 0.17 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "scale": 0.3, "sinStep": 0.11, "cosStep": 0.37 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [64, 128], "tolerance": 0.00002, "relTolerance": 0.0002 } }
+    },
+    {
+      "name": "sgmat_aligned_nobias_f32",
+      "provenance": {
+        "notes": "The no-bias arm of the same gate: sgmat renders its store without the bias term, which is a separate set of write sites from the bias arm."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [96, 96],
+          "data": { "kind": "fillFloat32", "scale": 0.7, "sinStep": 0.13, "cosStep": 0.29 }
+        },
+        "W": {
+          "dtype": "float32",
+          "shape": [96, 192],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.21, "cosStep": 0.17 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [96, 192], "tolerance": 0.00002, "relTolerance": 0.0002 } }
+    },
+    {
+      "name": "sgmat_aligned_bias_f16",
+      "provenance": {
+        "notes": "f16 operands drive the matrix units at their f16 component type while the accumulator and the fused gelu stay f32. The f16 gate admits M >= 2, so this also covers a row count under one tile."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.7, "sinStep": 0.13, "cosStep": 0.29 }
+        },
+        "W": {
+          "dtype": "float16",
+          "shape": [64, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.21, "cosStep": 0.17 }
+        },
+        "bias": {
+          "dtype": "float16",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "scale": 0.3, "sinStep": 0.11, "cosStep": 0.37 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [32, 128], "tolerance": 0.02, "relTolerance": 0.02 } }
+    }
+  ]
+}