sync 91d990483a17

Browse files

Files changed (7) hide show

README.md +28 -11
build/webgpu/bench.json +68 -1
build/webgpu/gemm-fast-gelu.wgsl.jinja +16 -61
build/webgpu/gemm-subgroup-matrix.wgsl.jinja +215 -59
build/webgpu/manifest.json +335 -153
build/webgpu/metadata.json +22 -8
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -18,17 +18,17 @@ See the [ONNX Runtime `GemmFastGelu` contrib-operator spec](https://github.com/m
 ## 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
@@ -36,13 +36,26 @@ See the [ONNX Runtime `GemmFastGelu` contrib-operator spec](https://github.com/m
 | --- | --- |
 | `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
@@ -51,10 +64,14 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `X` | `T` | — | — | Left operand of rank 2 or greater with shape `(..., K)`; every leading-axis coordinate identifies a row of the product. | required |
+| `W` | `T` | `2` | — | Right operand with shape `(K, N)`. | required |
+| `bias` | `T` | `1` | — | Optional bias with shape `(N)`, added before the activation. | optional |
 ## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `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
 | --- | --- |
 | `T` | `float32`, `float16` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `sgmat_direct_bias` — Uses supported subgroup matrices with direct activation loads and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_direct` — Uses supported subgroup matrices with direct activation loads and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_direct_bias_f16` — Uses supported subgroup matrices with direct activation loads and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_direct_f16` — Uses supported subgroup matrices with direct activation loads and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_bias` — Uses supported subgroup matrices with staged operands. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat` — Uses supported subgroup matrices with staged operands. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_bias_f16` — Uses supported subgroup matrices with staged operands. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_f16` — Uses supported subgroup matrices with staged operands. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 ## Device requirements
 Some implementation variants require `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.GemmFastGelu",
   "cases": [
     {
       "name": "gemmfastgelu-bert-base-b8-s384-h768-i3072",
@@ -47,6 +46,74 @@
       },
       "outputs": { "Y": { "shape": [8, 384, 3072], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * 3072 * 768 * 3072" }] }
     }
   ]
 }

 {
   "cases": [
     {
       "name": "gemmfastgelu-bert-base-b8-s384-h768-i3072",
       },
       "outputs": { "Y": { "shape": [8, 384, 3072], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * 3072 * 768 * 3072" }] }
+    },
+    {
+      "name": "gemmfastgelu-f32-short-m1-h768-i3072-sgmat-floor-pathology",
+      "preset": "model",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "X": { "shape": [1, 768], "dtype": "float32", "dist": "normal", "seed": 5250, "scale": 1 },
+        "W": { "shape": [768, 3072], "dtype": "float32", "dist": "normal", "seed": 5251, "scale": 1 },
+        "bias": { "shape": [3072], "dtype": "float32", "dist": "normal", "seed": 5252, "scale": 1 }
+      },
+      "outputs": { "Y": { "shape": [1, 3072], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 1 * 768 * 3072" }] }
+    },
+    {
+      "name": "gemmfastgelu-f32-short-m64-h768-i3072-control",
+      "preset": "model",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "X": { "shape": [64, 768], "dtype": "float32", "dist": "normal", "seed": 5270, "scale": 1 },
+        "W": { "shape": [768, 3072], "dtype": "float32", "dist": "normal", "seed": 5271, "scale": 1 },
+        "bias": { "shape": [3072], "dtype": "float32", "dist": "normal", "seed": 5272, "scale": 1 }
+      },
+      "outputs": { "Y": { "shape": [64, 3072], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 64 * 768 * 3072" }] }
+    },
+    {
+      "name": "gemmfastgelu-float32-bias-m128-k512-n96",
+      "preset": "model",
+      "inputs": {
+        "X": { "shape": [128, 512], "dtype": "float32", "dist": "normal", "seed": 7100, "scale": 0.3 },
+        "W": { "shape": [512, 96], "dtype": "float32", "dist": "normal", "seed": 7101, "scale": 0.2 },
+        "bias": { "shape": [96], "dtype": "float32", "dist": "normal", "seed": 7102, "scale": 0.3 }
+      },
+      "outputs": { "Y": { "shape": [128, 96], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": 12582912 }] }
+    },
+    {
+      "name": "gemmfastgelu-float32-bias-m128-k1024-n1024",
+      "preset": "model",
+      "inputs": {
+        "X": { "shape": [128, 1024], "dtype": "float32", "dist": "normal", "seed": 7100, "scale": 0.3 },
+        "W": { "shape": [1024, 1024], "dtype": "float32", "dist": "normal", "seed": 7101, "scale": 0.2 },
+        "bias": { "shape": [1024], "dtype": "float32", "dist": "normal", "seed": 7102, "scale": 0.3 }
+      },
+      "outputs": { "Y": { "shape": [128, 1024], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": 268435456 }] }
+    },
+    {
+      "name": "gemmfastgelu-float16-bias-m128-k512-n96",
+      "preset": "model",
+      "inputs": {
+        "X": { "shape": [128, 512], "dtype": "float16", "dist": "normal", "seed": 7100, "scale": 0.3 },
+        "W": { "shape": [512, 96], "dtype": "float16", "dist": "normal", "seed": 7101, "scale": 0.2 },
+        "bias": { "shape": [96], "dtype": "float16", "dist": "normal", "seed": 7102, "scale": 0.3 }
+      },
+      "outputs": { "Y": { "shape": [128, 96], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": 12582912 }] }
+    },
+    {
+      "name": "gemmfastgelu-float16-bias-m128-k1024-n1024",
+      "preset": "model",
+      "inputs": {
+        "X": { "shape": [128, 1024], "dtype": "float16", "dist": "normal", "seed": 7100, "scale": 0.3 },
+        "W": { "shape": [1024, 1024], "dtype": "float16", "dist": "normal", "seed": 7101, "scale": 0.2 },
+        "bias": { "shape": [1024], "dtype": "float16", "dist": "normal", "seed": 7102, "scale": 0.3 }
+      },
+      "outputs": { "Y": { "shape": [128, 1024], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": 268435456 }] }
     }
   ]
 }

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

@@ -10,7 +10,6 @@
 // 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; }
@@ -20,33 +19,6 @@ fn tanh_safe(x: f32) -> f32 {
 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" %}
-  // Clamp first: the hardware tanh evaluates (e^2v - 1)/(e^2v + 1), which is
-  // Inf/Inf = NaN once e^2v overflows (v > ~44). A GEMM accumulator reaches that
-  // easily because it grows with K. tanh is already +/-1 to full f32 precision by
-  // |v| ~ 9, so the clamp changes no representable result.
-  return tanh(clamp(v, -10.0, 10.0));
-{% 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;
@@ -89,71 +61,54 @@ fn main(
         // 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();
       }

 // 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" %}
 fn tanh_safe(x: f32) -> f32 {
   if (x > 10.0) { return 1.0; }
   if (x < -10.0) { return -1.0; }
 fn gelu_tanh(v: f32) -> f32 {
   return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
 }
 {% set epiOpen = "gelu_tanh(" if gemmEpi == "fastgelu" else ("fused_act(" if gemmEpi == "activation" else "") %}
 {% set epiClose = ")" if gemmEpi != "none" else "" %}
 const BK: u32 = 16u;
         // 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) {
           // 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;
           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) {
               aWord[{{ component }}u] = {{ tileT }}(a[am * K + ak + {{ component }}u]);
             }
 {% endfor %}
           }
           tileA[ar][ac4] = aWord;
         }
         for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + 256u) {
           // B stored [K, N]: logical B[k,n] = b[k*N + n]. Coalesce on n (contiguous).
           let bc = idx % BN;
           let br4 = idx / BN;
           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) {
               bWord[{{ component }}u] = {{ tileT }}(b[(bk + {{ component }}u) * N + bn]);
             }
 {% endfor %}
           }
           tileB[bc][br4] = bWord;
         }
         workgroupBarrier();
+{% set regT = "f32" %}{% set regCast = 1 if tileT != "f32" else 0 %}
+{% filter indent(8, true) %}
+let aRow = lid.y * TM;
+let bCol = lid.x * TN;
+for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
+  var av: array<vec4<{{ regT }}>, TM>;
+  var bv: array<vec4<{{ regT }}>, TN>;
+  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 %}; }
+  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 %}; }
+  for (var i: u32 = 0u; i < TM; i = i + 1u) {
+    for (var j: u32 = 0u; j < TN; j = j + 1u) {
+      acc[i * TN + j] = acc[i * TN + j] + dot(av[i], bv[j]);
+    }
+  }
+}
+{% endfilter %}
         workgroupBarrier();
       }

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

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
@@ -8,6 +5,7 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = mmaScalar %}
@@ -20,19 +18,52 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 {% 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; }
@@ -42,33 +73,6 @@ fn tanh_safe(x: f32) -> f32 {
 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" %}
-  // Clamp first: the hardware tanh evaluates (e^2v - 1)/(e^2v + 1), which is
-  // Inf/Inf = NaN once e^2v overflows (v > ~44). A GEMM accumulator reaches that
-  // easily because it grows with K. tanh is already +/-1 to full f32 precision by
-  // |v| ~ 9, so the clamp changes no representable result.
-  return tanh(clamp(v, -10.0, 10.0));
-{% 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") %}
@@ -79,13 +83,16 @@ 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
@@ -104,8 +111,8 @@ fn loadSHMA(tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   }
 }
 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;
@@ -121,32 +128,18 @@ fn loadSHMB(tile_base: u32, k_idx: u32, row: u32, c_idx: u32, bank_offset: u32)
       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;
-    {{ OUT }}[offset + row * OUT_F + col]  = {{ OUT_SCALAR }}({{ epiOpen }}scratch[src_slot][row * 8u + col]{{ BIAS_1 }}{{ epiClose }});
-    {{ OUT }}[offset + row * OUT_F + col2] = {{ OUT_SCALAR }}({{ epiOpen }}scratch[src_slot][row * 8u + col2]{{ BIAS_2 }}{{ epiClose }});
   }
 }
@@ -166,35 +159,179 @@ fn main(
   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>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, 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 %}
@@ -206,14 +343,18 @@ fn main(
 {% 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;
@@ -222,13 +363,28 @@ fn main(
   // 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<row_major>(&scratch[bank + {{ m * bTilesPerSubgroup + n }}u], 0u, matC{{ m }}{{ n }}, 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) %}
@@ -240,7 +396,7 @@ fn main(
     row,
     col,
     bank + {{ m * bTilesPerSubgroup + n }}u,
-    row_limit_{{ m }}
   );
 {% endfor %}
 {% endfor %}

 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = mmaScalar %}
 {% 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 useDirectStore = directStore is defined and directStore and gemmEpi == "none" %}
+{% set useDirectAStorage = directAStorage is defined and directAStorage %}
 {% set useDoubleBufferedB = doubleBufferedB is defined and doubleBufferedB %}
+{% set useBalancedAccumulation = balancedMatrixAccumulation is defined and balancedMatrixAccumulation and useDoubleBufferedB and not useDirectStore and operandScalar == "f32" %}
+{% set aTransposed = aTransposedStorage if aTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
 {% 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.
+{% if tileMValue > 32 %}
+{% set subtileCols = tileNValue %}
+{% set subtileRows = (tileMValue / 4)|int %}
+{% else %}
 {% set subtileCols = (tileNValue / 2)|int %}
 {% set subtileRows = 16 %}
+{% endif %}
 {% set aTilesPerSubgroup = (subtileRows / 8)|int %}
 {% set bTilesPerSubgroup = (subtileCols / 8)|int %}
 {% set subgroupCount = 4 %}
 {% set scratchBanks = subgroupCount * aTilesPerSubgroup * bTilesPerSubgroup %}
+{% macro matrixAccumulators() %}
+{% for m in range(aTilesPerSubgroup) %}
+{% for n in range(bTilesPerSubgroup) %}
+  var matC{{ m }}{{ n }}: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+{% if useBalancedAccumulation %}
+  var matD{{ m }}{{ n }}: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matE{{ m }}{{ n }}: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matF{{ m }}{{ n }}: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+{% endif %}
+{% endfor %}
+{% endfor %}
+{% endmacro %}
+{% macro storeMatrixAccumulators() %}
+{% for m in range(aTilesPerSubgroup) %}
+{% for n in range(bTilesPerSubgroup) %}
+  subgroupMatrixStore<row_major>(&scratch[bank + {{ m * bTilesPerSubgroup + n }}u], 0u, matC{{ m }}{{ n }}, 8u);
+{% if useBalancedAccumulation %}
+  subgroupMatrixStore<row_major>(&tile_B, (bank + {{ m * bTilesPerSubgroup + n }}u) * 64u, matD{{ m }}{{ n }}, 8u);
+{% endif %}
+{% endfor %}
+{% endfor %}
+{% endmacro %}
 fn tanh_safe(x: f32) -> f32 {
   if (x > 10.0) { return 1.0; }
   if (x < -10.0) { return -1.0; }
 fn gelu_tanh(v: f32) -> f32 {
   return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
 }
 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 SUB_ROWS:   u32 = {{ subtileRows }}u;
 const COLUMN_TILE_BASE: u32 = {{ columnTileBaseValue }}u;
 const ROW_TILE_BASE: u32 = {{ rowTileBaseValue }}u;
+{% if not useDirectAStorage %}
 var<workgroup> tile_A: array<{{ operandScalar }}, {{ tileMValue }} * {{ tileKValue }}>;
+{% endif %}
 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 }}>;
+{% if not useDirectAStorage %}
 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
   }
 }
+{% endif %}
 fn loadSHMB(tile_base: u32, k_idx: u32, row: u32, c_idx: u32, bank_offset: u32) {
   // 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;
       tile_B[bank_offset + k_local * TILE_COLS + n_local] = w[global_k * OUT_F + global_n];
     }
   }
 }
 {% 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 "" %}
+{% set ACC_1 = "(prior.x + (scratch[src_slot][row * 8u + col] + tile_B[src_slot * 64u + row * 8u + col]))" if useBalancedAccumulation else "scratch[src_slot][row * 8u + col]" %}
+{% set ACC_2 = "(prior.y + (scratch[src_slot][row * 8u + col2] + tile_B[src_slot * 64u + row * 8u + col2]))" if useBalancedAccumulation else "scratch[src_slot][row * 8u + col2]" %}
+fn storeOutput(offset: u32{% if needsColBase %}, col_base: u32{% endif %}, row: u32, col: u32, src_slot: u32, row_limit: i32{% if useBalancedAccumulation %}, prior: vec2<f32>{% endif %}) {
   if (row_limit > 0 && row < u32(row_limit)) {
     let col2: u32 = col + 1u;
+    {{ OUT }}[offset + row * OUT_F + col]  = {{ OUT_SCALAR }}({{ epiOpen }}{{ ACC_1 }}{{ BIAS_1 }}{{ epiClose }});
+    {{ OUT }}[offset + row * OUT_F + col2] = {{ OUT_SCALAR }}({{ epiOpen }}{{ ACC_2 }}{{ BIAS_2 }}{{ epiClose }});
   }
 }
   let w_global_base: u32 = (COLUMN_TILE_BASE + workgroup_id.x) * TILE_COLS;
   let subtile_id: u32 = local_idx / sg_size;
+{% if tileMValue > 32 %}
+  let subtile_idx: u32 = 0u;
+  let subtile_idy: u32 = subtile_id;
+{% else %}
   let subtile_idx: u32 = subtile_id / 2u;
   let subtile_idy: u32 = subtile_id % 2u;
+{% endif %}
   let base_A: u32 = subtile_idy * SUB_ROWS;
   let base_B: u32 = subtile_idx * SUB_COLS;
+{{ matrixAccumulators() }}
+{% if useDoubleBufferedB %}
+  // Large aligned f32 tiles do not benefit from staging A: each row band is
+  // consumed by exactly one subgroup. Load those matrices directly from x and
+  // ping-pong only B, which is reused by all four subgroups. The barrier that
+  // publishes the next B bank also proves the previous bank is no longer read,
+  // reducing the steady-state schedule from two barriers per K tile to one.
+  var b_bank: u32 = 0u;
+  loadSHMB(w_global_base, 0u, local_idx / {{ (tileKValue / 16)|int }}u, local_idx % {{ (tileKValue / 16)|int }}u, 0u);
+  workgroupBarrier();
+  for (var kidx: u32 = 0u; kidx < {{ kLoopEndValue }}; kidx += TILE_K) {
+    let b_bank_offset = b_bank * TILE_COLS * TILE_K;
+    for (var step: u32 = 0u; step < TILE_K; step += {{ 32 if useBalancedAccumulation else 8 }}u) {
+{% if useBalancedAccumulation %}
+{% for stream in range(4) %}
+    {
+{% set matrixAccumulator = ["matC", "matD", "matE", "matF"][stream] %}
+{% set matrixStepOffset = stream * 8 %}
+{% set dynamicATiles = aTilesPerSubgroup if aTilesPerSubgroup is defined else 2 %}
+{% set matrixAccumulator = matrixAccumulator | default("matC") %}
+{% set matrixStepOffset = matrixStepOffset | default(0) %}
+{% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
+{% set aTransposed = aTransposedStorage if aTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
+{% set bTransposed = bTransposedStorage if bTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
+{% set directA = directAStorage is defined and directAStorage %}
+{% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
+{% if directA %}
+      // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
+      // loads need no guarded workgroup staging.
+      let matrix_a_offset = (a_global_base + base_A) * IN_F + kidx + {{ matrixStep }};
+{% elif aTransposed %}
+      let matrix_a_offset = {{ matrixStep }} * TILE_ROWS + subtile_idy * SUB_ROWS;
+{% else %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
+{% endif %}
+{% for m in range(dynamicATiles) %}
+{% if directA %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&x, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * IN_F{% endif %}, IN_F);
+{% elif aTransposed %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, col_major>(&tile_A, matrix_a_offset + {{ m * 8 }}u, TILE_ROWS);
+{% else %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
+{% endif %}
+{% endfor %}
+{% if bTransposed %}
+      let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
+{% else %}
+      let matrix_b_offset = {{ B_BANK }}subtile_idx * SUB_COLS * TILE_K + {{ matrixStep }};
+{% endif %}
 {% for n in range(bTilesPerSubgroup) %}
+{% if bTransposed %}
+      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
+{% else %}
+      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u * TILE_K, TILE_K);
+{% endif %}
 {% endfor %}
+{% for m in range(dynamicATiles) %}
+{% for n in range(bTilesPerSubgroup) %}
+      {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
+{% endfor %}
+{% endfor %}
+    }
+{% endfor %}
+{% else %}
+{% set dynamicATiles = aTilesPerSubgroup if aTilesPerSubgroup is defined else 2 %}
+{% set matrixAccumulator = matrixAccumulator | default("matC") %}
+{% set matrixStepOffset = matrixStepOffset | default(0) %}
+{% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
+{% set aTransposed = aTransposedStorage if aTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
+{% set bTransposed = bTransposedStorage if bTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
+{% set directA = directAStorage is defined and directAStorage %}
+{% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
+{% if directA %}
+      // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
+      // loads need no guarded workgroup staging.
+      let matrix_a_offset = (a_global_base + base_A) * IN_F + kidx + {{ matrixStep }};
+{% elif aTransposed %}
+      let matrix_a_offset = {{ matrixStep }} * TILE_ROWS + subtile_idy * SUB_ROWS;
+{% else %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
+{% endif %}
+{% for m in range(dynamicATiles) %}
+{% if directA %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&x, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * IN_F{% endif %}, IN_F);
+{% elif aTransposed %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, col_major>(&tile_A, matrix_a_offset + {{ m * 8 }}u, TILE_ROWS);
+{% else %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
+{% endif %}
+{% endfor %}
+{% if bTransposed %}
+      let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
+{% else %}
+      let matrix_b_offset = {{ B_BANK }}subtile_idx * SUB_COLS * TILE_K + {{ matrixStep }};
+{% endif %}
+{% for n in range(bTilesPerSubgroup) %}
+{% if bTransposed %}
+      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
+{% else %}
+      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u * TILE_K, TILE_K);
+{% endif %}
 {% endfor %}
+{% for m in range(dynamicATiles) %}
+{% for n in range(bTilesPerSubgroup) %}
+      {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
+{% endfor %}
+{% endfor %}
+{% endif %}
+    }
+    let next_kidx = kidx + TILE_K;
+    if (next_kidx < {{ kLoopEndValue }}) {
+      let next_bank = 1u - b_bank;
+      loadSHMB(w_global_base, next_kidx, local_idx / {{ (tileKValue / 16)|int }}u, local_idx % {{ (tileKValue / 16)|int }}u, next_bank * TILE_COLS * TILE_K);
+      workgroupBarrier();
+      b_bank = next_bank;
+    }
+  }
+{% else %}
   for (var kidx: u32 = 0u; kidx < {{ kLoopEndValue }}; kidx += TILE_K) {
+{% if not useDirectAStorage %}
     loadSHMA(a_global_base, kidx, local_idx / {{ (tileKValue / 8)|int }}u, local_idx % {{ (tileKValue / 8)|int }}u);
+{% endif %}
     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 matrixAccumulator = matrixAccumulator | default("matC") %}
+{% set matrixStepOffset = matrixStepOffset | default(0) %}
+{% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
+{% set aTransposed = aTransposedStorage if aTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
 {% set bTransposed = bTransposedStorage if bTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
+{% set directA = directAStorage is defined and directAStorage %}
 {% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
+{% if directA %}
+      // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
+      // loads need no guarded workgroup staging.
+      let matrix_a_offset = (a_global_base + base_A) * IN_F + kidx + {{ matrixStep }};
+{% elif aTransposed %}
+      let matrix_a_offset = {{ matrixStep }} * TILE_ROWS + subtile_idy * SUB_ROWS;
+{% else %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
+{% endif %}
 {% for m in range(dynamicATiles) %}
+{% if directA %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&x, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * IN_F{% endif %}, IN_F);
+{% elif aTransposed %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, col_major>(&tile_A, matrix_a_offset + {{ m * 8 }}u, TILE_ROWS);
+{% else %}
       var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
+{% endif %}
 {% endfor %}
 {% if bTransposed %}
+      let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
 {% else %}
+      let matrix_b_offset = {{ B_BANK }}subtile_idx * SUB_COLS * TILE_K + {{ matrixStep }};
 {% endif %}
 {% for n in range(bTilesPerSubgroup) %}
 {% if bTransposed %}
 {% for m in range(dynamicATiles) %}
 {% for n in range(bTilesPerSubgroup) %}
+      {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
 {% endfor %}
 {% endfor %}
     }
     workgroupBarrier();
   }
+{% endif %}
+{% if useBalancedAccumulation %}
+  workgroupBarrier();
+{% endif %}
   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;
   // 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;
+{{ storeMatrixAccumulators() }}
+  workgroupBarrier();
+{% if useBalancedAccumulation %}
+  var pairTotals: array<vec2<f32>, {{ aTilesPerSubgroup * bTilesPerSubgroup }}>;
 {% for m in range(aTilesPerSubgroup) %}
 {% for n in range(bTilesPerSubgroup) %}
+  pairTotals[{{ m * bTilesPerSubgroup + n }}] = vec2<f32>(
+    scratch[bank + {{ m * bTilesPerSubgroup + n }}u][row * 8u + col] + tile_B[(bank + {{ m * bTilesPerSubgroup + n }}u) * 64u + row * 8u + col],
+    scratch[bank + {{ m * bTilesPerSubgroup + n }}u][row * 8u + col + 1u] + tile_B[(bank + {{ m * bTilesPerSubgroup + n }}u) * 64u + row * 8u + col + 1u]);
+{% endfor %}
+{% endfor %}
+  workgroupBarrier();
+{% for m in range(aTilesPerSubgroup) %}
+{% for n in range(bTilesPerSubgroup) %}
+  subgroupMatrixStore<row_major>(&scratch[bank + {{ m * bTilesPerSubgroup + n }}u], 0u, matE{{ m }}{{ n }}, 8u);
+  subgroupMatrixStore<row_major>(&tile_B, (bank + {{ m * bTilesPerSubgroup + n }}u) * 64u, matF{{ m }}{{ n }}, 8u);
 {% endfor %}
 {% endfor %}
   workgroupBarrier();
+{% endif %}
 {% 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) %}
     row,
     col,
     bank + {{ m * bTilesPerSubgroup + n }}u,
+    row_limit_{{ m }}{{ ", pairTotals[" ~ (m * bTilesPerSubgroup + n) ~ "]" if useBalancedAccumulation else "" }}
   );
 {% endfor %}
 {% endfor %}

build/webgpu/manifest.json CHANGED Viewed

@@ -2,38 +2,13 @@
   "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",
@@ -52,157 +27,336 @@
     "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)" }
         }
       ]
     },
@@ -210,16 +364,31 @@
       "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)"
           }
         }
       ]
@@ -228,16 +397,29 @@
       "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)"
           }
         }
       ]

   "domain": "com.microsoft",
   "name": "GemmFastGelu",
   "sinceVersion": 1,
+  "inputs": {
+    "X": { "dtype": "T" },
+    "W": { "dtype": "T", "rank": 2 },
+    "bias": { "dtype": "T", "rank": 1, "optional": true }
   },
+  "outputs": { "Y": { "dtype": "T", "rank": "ranks.X", "shape": "matmulShape(shapes.X, shapes.W)" } },
+  "typeConstraints": { "T": ["float32", "float16"] },
   "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",
     "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) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "sgmatContract": "wave32Effective and sgmatLayoutOk and sgmatResourcesFit",
     "usesF16": "tensorDtypes.X == \"float16\"",
     "transA": false,
     "transB": false,
+    "gemmEpilogue": "\"fastgelu\"",
+    "sgmatDirectRows": 64,
+    "sgmatDirectCols": 32,
+    "sgmatDirectStorageBytes": "2 * sgmatDirectCols * 32 * sgmatOperandBytes + sgmatDirectRows * sgmatDirectCols * 4",
+    "sgmatDirectContract": "wave32Effective and gfgRows > 0 and gfgRows % sgmatDirectRows == 0 and kDim % 32 == 0 and nDim % sgmatDirectCols == 0 and deviceWorkgroupCap >= 128 and sgmatDirectStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(gfgRows, sgmatDirectRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, sgmatDirectCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"
   },
+  "bindings": {
+    "x": { "arg": "X", "buffer": "read-only-storage", "elementType": "$xScalar" },
+    "w": { "arg": "W", "buffer": "read-only-storage", "elementType": "$wScalar" },
+    "y": { "arg": "Y", "buffer": "storage", "elementType": "$yScalar" },
+    "params": { "buffer": "uniform", "struct": [{ "name": "M", "type": "u32", "value": "gfgRows" }] },
+    "a": { "arg": "X", "buffer": "read-only-storage", "elementType": "$T" },
+    "b": { "arg": "W", "buffer": "read-only-storage", "elementType": "$T" },
+    "y_2": { "arg": "Y", "name": "y", "buffer": "storage", "elementType": "$T" }
   },
   "variants": [
+    {
+      "id": "sgmat_direct_bias",
+      "priority": 110,
+      "when": ["biasContract", "sgmatDirectContract", "tensorDtypes.X == \"float32\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }]
+      },
+      "derive": {
+        "hasBias": true,
+        "bTransposedStorage": true,
+        "xScalar": "\"f32\"",
+        "wScalar": "\"f32\"",
+        "bScalar": "\"f32\"",
+        "yScalar": "\"f32\"",
+        "mmaScalar": "\"f32\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatDirectRows",
+        "tileN": "sgmatDirectCols",
+        "tileK": 32,
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": true
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.sgmat_direct_bias",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "sgmat_direct",
+      "priority": 110,
+      "when": ["noBiasContract", "sgmatDirectContract", "tensorDtypes.X == \"float32\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }]
+      },
+      "derive": {
+        "hasBias": false,
+        "bTransposedStorage": true,
+        "xScalar": "\"f32\"",
+        "wScalar": "\"f32\"",
+        "bScalar": "\"f32\"",
+        "yScalar": "\"f32\"",
+        "mmaScalar": "\"f32\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatDirectRows",
+        "tileN": "sgmatDirectCols",
+        "tileK": 32,
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": true
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.sgmat_direct",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "sgmat_direct_bias_f16",
+      "priority": 110,
+      "when": ["biasContract", "sgmatDirectContract", "tensorDtypes.X == \"float16\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "hasBias": true,
+        "bTransposedStorage": true,
+        "xScalar": "\"f16\"",
+        "wScalar": "\"f16\"",
+        "bScalar": "\"f16\"",
+        "yScalar": "\"f16\"",
+        "mmaScalar": "\"f16\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatDirectRows",
+        "tileN": "sgmatDirectCols",
+        "tileK": 32,
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": false
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.sgmat_direct_bias_f16",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "sgmat_direct_f16",
+      "priority": 110,
+      "when": ["noBiasContract", "sgmatDirectContract", "tensorDtypes.X == \"float16\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "hasBias": false,
+        "bTransposedStorage": true,
+        "xScalar": "\"f16\"",
+        "wScalar": "\"f16\"",
+        "bScalar": "\"f16\"",
+        "yScalar": "\"f16\"",
+        "mmaScalar": "\"f16\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatDirectRows",
+        "tileN": "sgmatDirectCols",
+        "tileK": 32,
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": false
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.sgmat_direct_f16",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
     {
       "id": "sgmat_bias",
       "priority": 100,
+      "when": ["biasContract", "sgmatContract", "tensorDtypes.X == \"float32\""],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }]
       },
+      "derive": {
         "hasBias": true,
         "bTransposedStorage": true,
+        "xScalar": "\"f32\"",
+        "wScalar": "\"f32\"",
+        "bScalar": "\"f32\"",
+        "yScalar": "\"f32\"",
+        "mmaScalar": "\"f32\"",
         "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": 32,
+        "tileN": 64,
+        "tileK": 32,
+        "directAStorage": false,
+        "doubleBufferedB": false,
+        "balancedMatrixAccumulation": false
       },
       "passes": [
         {
           "id": "main",
+          "name": "GemmFastGelu.sgmat_bias",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
         }
       ]
     },
     {
       "id": "sgmat",
       "priority": 100,
+      "when": ["noBiasContract", "sgmatContract", "tensorDtypes.X == \"float32\""],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }]
       },
+      "derive": {
         "hasBias": false,
         "bTransposedStorage": true,
+        "xScalar": "\"f32\"",
+        "wScalar": "\"f32\"",
+        "bScalar": "\"f32\"",
+        "yScalar": "\"f32\"",
+        "mmaScalar": "\"f32\"",
         "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": 32,
+        "tileN": 64,
+        "tileK": 32,
+        "directAStorage": false,
+        "doubleBufferedB": false,
+        "balancedMatrixAccumulation": false
       },
       "passes": [
         {
           "id": "main",
+          "name": "GemmFastGelu.sgmat",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "sgmat_bias_f16",
+      "priority": 100,
+      "when": ["biasContract", "sgmatContract", "tensorDtypes.X == \"float16\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "hasBias": true,
+        "bTransposedStorage": true,
+        "xScalar": "\"f16\"",
+        "wScalar": "\"f16\"",
+        "bScalar": "\"f16\"",
+        "yScalar": "\"f16\"",
+        "mmaScalar": "\"f16\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": 32,
+        "tileN": 64,
+        "tileK": 32,
+        "directAStorage": false,
+        "doubleBufferedB": false,
+        "balancedMatrixAccumulation": false
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.sgmat_bias_f16",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "sgmat_f16",
+      "priority": 100,
+      "when": ["noBiasContract", "sgmatContract", "tensorDtypes.X == \"float16\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "hasBias": false,
+        "bTransposedStorage": true,
+        "xScalar": "\"f16\"",
+        "wScalar": "\"f16\"",
+        "bScalar": "\"f16\"",
+        "yScalar": "\"f16\"",
+        "mmaScalar": "\"f16\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": 32,
+        "tileN": 64,
+        "tileK": 32,
+        "directAStorage": false,
+        "doubleBufferedB": false,
+        "balancedMatrixAccumulation": false
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.sgmat_f16",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
         }
       ]
     },
       "id": "tiled_bias",
       "priority": 0,
       "when": ["biasContract"],
+      "derive": { "hasBias": true, "rowBias": true, "scalarBias": false, "columnBias": false },
       "passes": [
         {
           "id": "main",
           "name": "GemmFastGelu.Bias",
           "shader": "gemm-fast-gelu.wgsl.jinja",
+          "bindings": [
+            "a",
+            "b",
+            "bias",
+            "y_2",
+            {
+              "name": "params",
+              "struct": [
+                { "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 }
+              ]
+            }
+          ],
           "dispatch": {
+            "x": "min(ceilDiv(nDim, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": "min(ceilDiv(gfgRows, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))"
           }
         }
       ]
       "id": "tiled",
       "priority": 0,
       "when": ["noBiasContract"],
+      "derive": { "hasBias": false },
       "passes": [
         {
           "id": "main",
           "name": "GemmFastGelu",
           "shader": "gemm-fast-gelu.wgsl.jinja",
+          "bindings": [
+            "a",
+            "b",
+            "y_2",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "M", "type": "u32", "value": "gfgRows" },
+                { "name": "N", "type": "u32", "value": "nDim" },
+                { "name": "K", "type": "u32", "value": "kDim" },
+                { "name": "alpha", "type": "f32", "value": 1 }
+              ]
+            }
+          ],
           "dispatch": {
+            "x": "min(ceilDiv(nDim, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": "min(ceilDiv(gfgRows, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))"
           }
         }
       ]

build/webgpu/metadata.json CHANGED Viewed

@@ -1,19 +1,33 @@
 {
   "name": "com.microsoft.GemmFastGelu",
-  "id": "_com_microsoft_gemmfastgelu_webgpu_720afad",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "CPFUVoaXFHE1nKK5oZMKmB0LohpzyltnfuRq/70nJAk=",
-      "gemm-fast-gelu.wgsl.jinja": "XL5rMA9oUlUJZCMHq2PZAzmORxFZoz/xa5Mo05/aK3k=",
-      "gemm-subgroup-matrix.wgsl.jinja": "JJZezmRCzycCZ5Y6RTmbisasCFLbVAgwkSkV4Sct7eY=",
-      "manifest.json": "bcpv0SoGGVwU23+3nF7YJ2BSX80r64AInCEZgAuCRoI=",
-      "test.json": "WNyAFlYnnkTg3N4r9ja8Kh2HwlO8jr7EncOBBa9oI48="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GemmFastGelu" }
 }

 {
   "name": "com.microsoft.GemmFastGelu",
+  "id": "_com_microsoft_gemmfastgelu_webgpu_a713abb",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "oqlnhXJEZ8273ZAZC7ET+BMVMM8J84BrfZAVpaboErU=",
+      "gemm-fast-gelu.wgsl.jinja": "qh+CurhpiRX3buxeBnrOX2J2mzPrlqMP+m++faB5SfI=",
+      "gemm-subgroup-matrix.wgsl.jinja": "aaVCZdPC5akSTqDdYvM/MP8NnncAcK9kwSbvrd4Bvj8=",
+      "manifest.json": "EespK6n0UbB6+9TwNvJ40Pp/lYYHsHCO0BCKXxLRZhM=",
+      "test.json": "LFwbPPwea77gPJaoCjDJtA/uWfmqDYMUtIoDxd8iphY="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "sgmat_direct_bias": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_direct": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_direct_bias_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_direct_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_bias": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_bias_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "tiled_bias": ["gemm-fast-gelu.wgsl.jinja"],
+      "tiled": ["gemm-fast-gelu.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/test.json CHANGED Viewed

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