sync 91d990483a17

Browse files

Files changed (12) hide show

README.md +28 -10
build/webgpu/bench.json +224 -3
build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja +12 -13
build/webgpu/manifest.json +333 -171
build/webgpu/matmul-band-vec4.wgsl.jinja +84 -0
build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja +32 -80
build/webgpu/matmul-tiled-general-reg.wgsl.jinja +66 -42
build/webgpu/matmul-tiled-general.wgsl.jinja +31 -34
build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja +9 -10
build/webgpu/metadata.json +31 -12
build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja +94 -6
build/webgpu/test.json +282 -28

README.md CHANGED Viewed

@@ -18,16 +18,16 @@ See the [ONNX Runtime `FusedMatMul` contrib-operator spec](https://github.com/mi
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `A` | `A` | `T` | — | — | N-dimensional matrix A. | required |
-| `B` | `B` | `T` | — | — | N-dimensional matrix B. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `Y` | `Y` | `T` | derived | derived; see description | Matrix-multiplication result whose shape follows NumPy-style rules after applying the requested batch and matrix transpositions. | required |
 ## Attributes
@@ -47,17 +47,31 @@ Default values (overridable per request):
 | --- | --- |
 | `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
 - [`fused-matmul-subgroup-matrix.wgsl.jinja`](build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja)
 - [`matmul-subgroup-matrix-ext.wgsl.jinja`](build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja)
 - [`matmul-tiled-general-reg.wgsl.jinja`](build/webgpu/matmul-tiled-general-reg.wgsl.jinja)
 - [`matmul-tiled-general.wgsl.jinja`](build/webgpu/matmul-tiled-general.wgsl.jinja)
@@ -66,10 +80,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 |
+| --- | --- | --- | --- | --- | --- |
+| `A` | `T` | — | — | N-dimensional matrix A. | required |
+| `B` | `T` | — | — | N-dimensional matrix B. | required |
 ## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `Y` | `T` | derived | derived | Matrix-multiplication result whose shape follows NumPy-style rules after applying the requested batch and matrix transpositions. | required |
 ## Attributes
 | --- | --- |
 | `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.
+- `subgroup_matrix_transbatch_b_f16` — Aligned rank-3 transBatchB product with sufficient reduction depth and output tiles to amortize subgroup-matrix staging.
+- `subgroup_matrix_transbatch_b_f32` — Aligned rank-3 transBatchB product with sufficient reduction depth and output tiles to amortize subgroup-matrix staging.
+- `rank2_band_vec4_splitk` — Splits the vec4 band's K axis across up to sixteen workgroups. Each range writes an f32 partial band with alpha applied, and a combine pass sums the partials.
+- `rank2_band_vec4` — Few-row band for a rank-2 product without transposes. Each lane owns one vec4 column group and one accumulator per row, so every B word feeds all 2 to 16 rows; alpha is applied at the store.
+- `rank2_band_vec4_f32_preferred` — Few-row band for a rank-2 product without transposes. Each lane owns one vec4 column group and one accumulator per row, so every B word feeds all 2 to 16 rows; alpha is applied at the store.
+- `subgroup_matrix_splitk` — Partitions the K reduction of small-M rank-2 products across subgroup-matrix workgroups, then combines float32 partials that already include alpha.
+- `plain_rank2_tiled_reg` — Register-blocked rank-2 `Y = alpha * A @ B` specialization for non-transposed inputs on tiers without subgroup-matrix support.
+- `transbatch_b_tiled_reg` — Register-blocked logical rank3 product with an interleaved physical B batch axis.
 ## 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
 - [`fused-matmul-subgroup-matrix.wgsl.jinja`](build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja)
+- [`matmul-band-vec4.wgsl.jinja`](build/webgpu/matmul-band-vec4.wgsl.jinja)
 - [`matmul-subgroup-matrix-ext.wgsl.jinja`](build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja)
 - [`matmul-tiled-general-reg.wgsl.jinja`](build/webgpu/matmul-tiled-general-reg.wgsl.jinja)
 - [`matmul-tiled-general.wgsl.jinja`](build/webgpu/matmul-tiled-general.wgsl.jinja)
 ## 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.FusedMatMul",
   "cases": [
     {
       "name": "fusedmatmul-f32-attn-scores-transB-8x512x64",
@@ -164,7 +163,7 @@
       "name": "fusedmatmul-f16-rank4-by-rank2-shared-weight-b2h8-m512-k2048-n512-pathology",
       "preset": "stress",
       "provenance": {
-        "notes": "Production-scale batched projection with a shared rank-2 weight. This common numpy.matmul broadcast shape is excluded from the subgroup-matrix rank combinations and locks the register-blocked portable fallback that replaced the lower-throughput generic tiled kernel."
       },
       "attrs": { "alpha": 1 },
       "vars": { "dtype": "float16", "M": 512, "K": 2048, "N": 512 },
@@ -191,7 +190,7 @@
       "name": "fusedmatmul-f16-transbatch-b-8x512x2048x512-pathology",
       "preset": "stress",
       "provenance": {
-        "notes": "The operator's valid transBatchB semantic at a production projection size. No subgroup-matrix or specialized trans-batch-B variant accepts this layout, so it selects the lower-throughput generic tiled kernel."
       },
       "attrs": { "alpha": 1, "transBatchB": 1 },
       "vars": { "batch": 8, "M": 512, "K": 2048, "N": 512 },
@@ -278,6 +277,228 @@
       },
       "outputs": { "Y": { "shape": [8, 4, 512], "dtype": "float16", "dist": "empty" } },
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 4 * 64 * 512" }] }
     }
   ]
 }

 {
   "cases": [
     {
       "name": "fusedmatmul-f32-attn-scores-transB-8x512x64",
       "name": "fusedmatmul-f16-rank4-by-rank2-shared-weight-b2h8-m512-k2048-n512-pathology",
       "preset": "stress",
       "provenance": {
+        "notes": "A production-scale batched projection shares one rank-2 weight across batches. This NumPy-style broadcast shape exercises the register-blocked portable route."
       },
       "attrs": { "alpha": 1 },
       "vars": { "dtype": "float16", "M": 512, "K": 2048, "N": 512 },
       "name": "fusedmatmul-f16-transbatch-b-8x512x2048x512-pathology",
       "preset": "stress",
       "provenance": {
+        "notes": "Interleaved B batches in a production-scale projection; compares stride-aware subgroup matrices with the portable tiled path."
       },
       "attrs": { "alpha": 1, "transBatchB": 1 },
       "vars": { "batch": 8, "M": 512, "K": 2048, "N": 512 },
       },
       "outputs": { "Y": { "shape": [8, 4, 512], "dtype": "float16", "dist": "empty" } },
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 4 * 64 * 512" }] }
+    },
+    {
+      "name": "fusedmatmul-f32-decode-gemv-m1-k4096-n4096-alpha-half",
+      "provenance": {
+        "notes": "A 1x4096 by 4096x4096 decode GEMV with alpha=0.5 exercises the single-row route with a baked non-unit output multiplier."
+      },
+      "preset": "stress",
+      "attrs": { "alpha": 0.5 },
+      "vars": { "dtype": "float32", "M": 1, "K": 4096, "N": 4096 },
+      "inputs": {
+        "A": { "shape": [1, 4096], "dtype": "float32", "dist": "normal", "seed": 740, "scale": 0.1 },
+        "B": { "shape": [4096, 4096], "dtype": "float32", "dist": "normal", "seed": 741, "scale": 0.05 }
+      },
+      "outputs": { "Y": { "shape": [1, 4096], "dtype": "float32", "dist": "empty" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(args.M * args.K + args.K * args.N + args.M * args.N) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedmatmul-f32-transbatch-b-8x512x2048x512-pathology",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Interleaved B batches in a production-scale projection; compares stride-aware subgroup matrices with the portable tiled path."
+      },
+      "attrs": { "alpha": 1, "transBatchB": 1 },
+      "vars": { "batch": 8, "M": 512, "K": 2048, "N": 512 },
+      "inputs": {
+        "A": { "shape": [8, 512, 2048], "dtype": "float32", "dist": "normal", "seed": 722, "scale": 0.05 },
+        "B": { "shape": [2048, 8, 512], "dtype": "float32", "dist": "normal", "seed": 723, "scale": 0.05 }
+      },
+      "outputs": { "Y": { "shape": [8, 512, 512], "dtype": "float32", "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 512 * 512 * 2048" }] }
+    },
+    {
+      "name": "fusedmatmul-float16-transbatch-b-k128-workgroups64",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Aligned interleaved B layout at the reduction-length and 64-matrix-workgroup selector floors."
+      },
+      "attrs": { "alpha": 0.5, "transBatchB": 1 },
+      "vars": { "batch": 2, "M": 128, "K": 128, "N": 512 },
+      "inputs": {
+        "A": { "shape": [2, 128, 128], "dtype": "float16", "dist": "normal", "seed": 722, "scale": 0.05 },
+        "B": { "shape": [128, 2, 512], "dtype": "float16", "dist": "normal", "seed": 723, "scale": 0.05 }
+      },
+      "outputs": { "Y": { "shape": [2, 128, 512], "dtype": "float16", "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * args.batch * args.M * args.N * args.K" }] }
+    },
+    {
+      "name": "fusedmatmul-float32-transbatch-b-k128-workgroups64",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Aligned interleaved B layout at the reduction-length and 64-matrix-workgroup selector floors."
+      },
+      "attrs": { "alpha": 0.5, "transBatchB": 1 },
+      "vars": { "batch": 2, "M": 128, "K": 128, "N": 512 },
+      "inputs": {
+        "A": { "shape": [2, 128, 128], "dtype": "float32", "dist": "normal", "seed": 722, "scale": 0.05 },
+        "B": { "shape": [128, 2, 512], "dtype": "float32", "dist": "normal", "seed": 723, "scale": 0.05 }
+      },
+      "outputs": { "Y": { "shape": [2, 128, 512], "dtype": "float32", "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * args.batch * args.M * args.N * args.K" }] }
+    },
+    {
+      "name": "transbatch_b_portable_reg_aligned_float16",
+      "attrs": { "transBatchB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [4, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [128, 4, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.07, "offset": 0.3 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 128, 512] } },
+      "provenance": {
+        "notes": "Portable register tile transBatchB: physical B[K,batch,N], natural tile count reaches the shared occupancy gate; negative alpha and positive inputs avoid cancellation."
+      },
+      "preset": "smoke",
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, 2)" }] }
+    },
+    {
+      "name": "transbatch_b_portable_reg_tails_float16",
+      "attrs": { "transBatchB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [3, 129, 131],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [131, 3, 513],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.07, "offset": 0.3 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [3, 129, 513] } },
+      "provenance": {
+        "notes": "Portable register tile transBatchB: physical B[K,batch,N], natural tile count reaches the shared occupancy gate; negative alpha and positive inputs avoid cancellation."
+      },
+      "preset": "smoke",
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, 2)" }] }
+    },
+    {
+      "name": "transbatch_b_portable_reg_aligned_float32",
+      "attrs": { "transBatchB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [4, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [128, 4, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.07, "offset": 0.3 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 128, 512] } },
+      "provenance": {
+        "notes": "Portable register tile transBatchB: physical B[K,batch,N], natural tile count reaches the shared occupancy gate; negative alpha and positive inputs avoid cancellation."
+      },
+      "preset": "smoke",
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, 2)" }] }
+    },
+    {
+      "name": "transbatch_b_portable_reg_tails_float32",
+      "attrs": { "transBatchB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 129, 131],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [131, 3, 513],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.07, "offset": 0.3 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [3, 129, 513] } },
+      "provenance": {
+        "notes": "Portable register tile transBatchB: physical B[K,batch,N], natural tile count reaches the shared occupancy gate; negative alpha and positive inputs avoid cancellation."
+      },
+      "preset": "smoke",
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, 2)" }] }
+    },
+    {
+      "name": "f32-band-preference-m4-k2048-n4096",
+      "preset": "model",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [4, 2048], "dist": "normal", "seed": 9214, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [2048, 4096], "dist": "normal", "seed": 9215, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 4096], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": 67108864 }] }
+    },
+    {
+      "name": "f32-band-preference-m4-k8192-n4096",
+      "preset": "model",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [4, 8192], "dist": "normal", "seed": 9216, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [8192, 4096], "dist": "normal", "seed": 9217, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 4096], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": 268435456 }] }
+    },
+    {
+      "name": "f32-band-preference-m8-k2048-n4096",
+      "preset": "model",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [8, 2048], "dist": "normal", "seed": 9222, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [2048, 4096], "dist": "normal", "seed": 9223, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 4096], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": 134217728 }] }
+    },
+    {
+      "name": "f32-band-preference-m8-k8192-n4096",
+      "preset": "model",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [8, 8192], "dist": "normal", "seed": 9224, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [8192, 4096], "dist": "normal", "seed": 9225, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 4096], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": 536870912 }] }
+    },
+    {
+      "name": "f32-band-preference-m16-k2048-n4096",
+      "preset": "model",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [16, 2048], "dist": "normal", "seed": 9230, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [2048, 4096], "dist": "normal", "seed": 9231, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [16, 4096], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": 268435456 }] }
+    },
+    {
+      "name": "f32-band-preference-m16-k8192-n4096",
+      "preset": "model",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [16, 8192], "dist": "normal", "seed": 9232, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [8192, 4096], "dist": "normal", "seed": 9233, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [16, 4096], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": 1073741824 }] }
     }
   ]
 }

build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja CHANGED Viewed

@@ -1,12 +1,9 @@
-// com.microsoft.FusedMatMul subgroup-matrix specialization: Y = alpha * A @ op(B),
-// op(B) = B^T when transB. Dense batches map through workgroup_id.z, and
-// M-tail rows are guarded by row_limit. transA uses the tiled specialization.
 // Alignment gates keep K % 32 == 0 and N % 64 == 0 so subgroupMatrixLoad never
 // sees partial 8x8 tiles. The batch is required to match between A and B
 // (no broadcast) because a_base/b_base both index by the same workgroup_id.z.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
@@ -14,6 +11,7 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = fScalar %}
@@ -34,7 +32,7 @@ const A_M_STRIDE: u32 = K;
 {% endif %}
 const B_BATCH_STRIDE: u32 = K * N;
 const C_BATCH_STRIDE: u32 = M * N;
-const ALPHA: {{ accScalar }} = {{ accScalar }}({{ source.alpha }});
 const TILE_COLS: u32 = 64u;
 const TILE_ROWS: u32 = 32u;
 const TILE_K: u32 = 32u;
@@ -127,15 +125,16 @@ fn main(
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
-      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
-      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
-      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
-      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
-      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
-      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);

+// com.microsoft.FusedMatMul subgroup-matrix specialization: Y = alpha * op(A) @ op(B).
+// transA and transB transpose the corresponding matrix operand on load.
+// Dense batches map through workgroup_id.z, and M-tail rows are guarded by row_limit.
 // Alignment gates keep K % 32 == 0 and N % 64 == 0 so subgroupMatrixLoad never
 // sees partial 8x8 tiles. The batch is required to match between A and B
 // (no broadcast) because a_base/b_base both index by the same workgroup_id.z.
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = fScalar %}
 {% endif %}
 const B_BATCH_STRIDE: u32 = K * N;
 const C_BATCH_STRIDE: u32 = M * N;
+const ALPHA: {{ accScalar }} = {{ accScalar }}({{ alpha }});
 const TILE_COLS: u32 = 64u;
 const TILE_ROWS: u32 = 32u;
 const TILE_K: u32 = 32u;
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set directInputs = directMatrixInputs is defined and directMatrixInputs %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+{% for r in range(2) %}
+      var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});
+{% endfor %}
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+{% for c in range(4) %}
+      var matB{{ c }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&{{ "xm" if directInputs else "tile_B" }}, matrix_b_offset{% if c > 0 %} + {{ c * 8 }}u * TILE_K{% endif %}, {{ "N" if directInputs else "TILE_K" }});
+{% endfor %}
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);

build/webgpu/manifest.json CHANGED Viewed

@@ -2,91 +2,43 @@
   "domain": "com.microsoft",
   "name": "FusedMatMul",
   "sinceVersion": 1,
-  "description": "Matrix product of two N-dimensional tensors `A` and `B`, following NumPy-style matrix-multiplication broadcasting. Supports optional transposition of either operand's last two dimensions, optional batch-dimension transposition, and a scalar `alpha` multiplier. Float32 and float16 are supported; double and bfloat16 are not.",
-  "inputs": [
-    { "role": "A", "dtype": "T", "description": "N-dimensional matrix A." },
-    { "role": "B", "dtype": "T", "description": "N-dimensional matrix B." }
-  ],
-  "outputs": [
-    {
-      "role": "Y",
       "dtype": "T",
       "rank": "max(ranks.A, ranks.B) - (1 if ranks.A == 1 or ranks.B == 1 else 0)",
-      "shape": "matmulShape(logicalAShape, logicalBShape)",
-      "description": "Matrix-multiplication result whose shape follows NumPy-style rules after applying the requested batch and matrix transpositions."
     }
-  ],
-  "attributes": { "alpha": 1, "transA": 0, "transB": 0, "transBatchA": 0, "transBatchB": 0 },
-  "attributeDescriptions": {
-    "alpha": "Scalar multiplier applied to the product of the input tensors.",
-    "transA": "When non-zero, transposes `A` on its last two dimensions before multiplication.",
-    "transB": "When non-zero, transposes `B` on its last two dimensions before multiplication.",
-    "transBatchA": "When non-zero, transposes `A` on its first dimension and batch dimensions (dim-1 to dim-rank-2) before multiplication.",
-    "transBatchB": "When non-zero, transposes `B` on its first dimension and batch dimensions (dim-1 to dim-rank-2) before multiplication."
   },
-  "typeConstraints": { "T": ["float32", "float16"] },
-  "args": {
-    "A": { "kind": "tensor", "semantic": "A", "role": "input" },
-    "B": { "kind": "tensor", "semantic": "B", "role": "input" },
-    "Y": { "kind": "tensor", "semantic": "Y", "role": "output" }
   },
   "tunables": {
-    "TILED_REG_MIN_WORKGROUPS": 64,
-    "GEMV_TARGET_BLOCKS": 512,
-    "SUBGROUP_MATRIX_MIN_M": 2,
-    "SUBGROUP_MATRIX_SPLITK_TARGET_WGS": 512,
-    "SUBGROUP_MATRIX_SPLITK_MIN_K": 1024,
-    "SUBGROUP_MATRIX_SPLITK_MAX_TILES": 128
-  },
-  "bindingSets": {
-    "matrix": [
-      { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "b", "arg": "B", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" }
-    ],
-    "sgmatSplitPartial": [
-      { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "b", "arg": "B", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" }
-    ],
-    "sgmatSplitCombine": [
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "numel(shapes.Y)" }] }
-      }
-    ],
-    "gemvF32": [
-      { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "b",
-        "arg": "B",
-        "semantic": "B",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "vec4<f32>"
-      },
-      { "name": "c", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "K", "type": "u32", "value": "dim(shapes.A, 1)" },
-            { "name": "N4", "type": "u32", "value": "dim(shapes.B, 1) / 4" }
-          ]
-        }
-      }
-    ]
   },
   "derive": {
     "gemvWorkgroups": "ceilDiv(dim(shapes.B, 1), 128)",
     "gemvSliceCap": "min(32, device.limits.maxComputeWorkgroupSizeY, floor(device.limits.maxComputeInvocationsPerWorkgroup / 32), floor(device.limits.maxComputeWorkgroupStorageSize / 512))",
-    "gemvSlices": "max(1, min(gemvSliceCap, max(8, pow2ceil(ceilDiv(tunables.GEMV_TARGET_BLOCKS, gemvWorkgroups)))))",
     "batchMovedAShape": "moveAxis(shapes.A, 0, -2) if attrs.transBatchA != 0 else shapes.A",
     "batchMovedBShape": "moveAxis(shapes.B, 0, -2) if attrs.transBatchB != 0 else shapes.B",
     "logicalAShape": "moveAxis(batchMovedAShape, -1, -2) if attrs.transA != 0 and ranks.A > 1 else batchMovedAShape",
@@ -107,28 +59,224 @@
     "sgmatSplitK8Ok": "sgmatSplitKDepth % 256 == 0",
     "sgmatSplitK4Ok": "sgmatSplitKDepth % 128 == 0",
     "sgmatSplitK2Ok": "sgmatSplitKDepth % 64 == 0",
-    "sgmatSplitK": "32 if (sgmatSplitKWant > 16 and sgmatSplitK32Ok) else (16 if (sgmatSplitKWant > 8 and sgmatSplitK16Ok) else (8 if (sgmatSplitKWant > 4 and sgmatSplitK8Ok) else (4 if (sgmatSplitKWant > 2 and sgmatSplitK4Ok) else (2 if sgmatSplitK2Ok else 1))))"
   },
   "variants": [
     {
       "id": "f32_m1_gemv_vec4",
       "priority": 30,
-      "when": ["dtypes.T == \"f32\"", "attrs.alpha == 1", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "ranks.A == 2", "ranks.B == 2", "ranks.Y == 2", "dim(shapes.A, 0) == 1", "dim(shapes.Y, 0) == 1", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 1) == dim(shapes.B, 1)", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "ceil(dim(shapes.B, 1) / 128) <= device.limits.maxComputeWorkgroupsPerDimension"],
-      "constants": { "T": "\"f32\"", "usesF16": false, "gemvSlices": "gemvSlices" },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.F32M1GemvVec4",
           "shader": "matmul-vector-matrix-vec4.wgsl.jinja",
-          "bindings": "gemvF32",
           "dispatch": { "x": "ceil(dim(shapes.B, 1) / 128)" }
         }
       ]
     },
     {
       "id": "subgroup_matrix_splitk",
-      "description": "Partitions the K reduction across workgroups while retaining subgroup-matrix tiles, then combines float32 partials that already carry alpha. It serves small-M rank-two products whose unsplit tile grids cannot occupy the device.",
       "priority": 12,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [
@@ -136,8 +284,7 @@
           { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
         ]
       },
-      "when": ["(dtypes.T == \"f16\" or dtypes.T == \"f32\") and f16Ok(dtypes.T)", "fusedSgmatRank2Ok", "dim(shapes.A, 0) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, 1) >= tunables.SUBGROUP_MATRIX_SPLITK_MIN_K", "dim(shapes.B, 1) % 64 == 0", "sgmatSplitK >= 2", "sgmatOutTiles < tunables.SUBGROUP_MATRIX_SPLITK_MAX_TILES", "sgmatSplitK * numel(shapes.Y) * 4 <= device.limits.maxStorageBufferBindingSize", "sgmatSplitK <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(dim(shapes.Y, 1), 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(dim(shapes.Y, 0), 32) <= device.limits.maxComputeWorkgroupsPerDimension", "subgroupMatrixResourcesFit", "wave32Effective"],
-      "constants": {
         "usesF16": "dtypes.T == \"f16\"",
         "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
         "scalar": "dtypes.T",
@@ -161,39 +308,39 @@
         {
           "id": "partial",
           "name": "FusedMatMul.SubgroupMatrixSplitK",
-          "source": {
-            "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
-            "inputs": {
-              "aShape": ["dim(shapes.A, 0)", "dim(shapes.A, 1)"],
-              "bShape": ["dim(shapes.B, 0)", "dim(shapes.B, 1)"],
-              "aRank": 2,
-              "bRank": 2
-            }
           },
-          "bindings": "sgmatSplitPartial",
           "dispatch": { "x": "ceilDiv(dim(shapes.Y, 1), 64)", "y": "ceilDiv(dim(shapes.Y, 0), 32)", "z": "sgmatSplitK" }
         },
         {
           "id": "combine",
           "name": "FusedMatMul.SubgroupMatrixSplitKCombine",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"", "intMode": false }
-          },
-          "bindings": "sgmatSplitCombine",
-          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": 256 }
         }
       ]
     },
     {
       "id": "subgroup_matrix_tail_broadcast",
       "priority": 11,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
       },
-      "when": ["dtypes.T == \"f16\"", "f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "((ranks.A == 2 and ranks.B == 2 and ranks.Y == 2) or (ranks.A == 4 and ranks.B == 3 and ranks.Y == 4 and dim(shapes.Y, 0) == dim(shapes.A, 0) and (dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1) and dim(shapes.Y, 1) == max(dim(shapes.A, 1), dim(shapes.B, 0))))", "dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 2)", "dim(shapes.A, ranks.A - 2) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, ranks.A - 1) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "dim(shapes.Y, ranks.Y - 2) == dim(shapes.A, ranks.A - 2)", "dim(shapes.Y, ranks.Y - 1) == dim(shapes.B, ranks.B - 1)", "ceil(dim(shapes.B, ranks.B - 1) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, ranks.A - 2) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"],
-      "constants": {
         "hasBias": false,
         "usesF16": true,
         "fScalar": "\"f16\"",
@@ -212,18 +359,17 @@
         {
           "id": "main",
           "name": "FusedMatMul.SubgroupMatrixTailBroadcast",
-          "source": {
-            "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
-            "inputs": { "aShape": "shapes.A", "bShape": "shapes.B", "aRank": "ranks.A", "bRank": "ranks.B" }
-          },
-          "bindings": "matrix",
-          "dispatch": { "x": "ceil(constants.N / 64)", "y": "ceil(constants.M / 32)", "z": "constants.batchCount" }
         }
       ]
     },
     {
       "id": "subgroup_matrix",
       "priority": 10,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [
@@ -231,8 +377,7 @@
           { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
         ]
       },
-      "when": ["f16Ok(dtypes.T)", "attrs.transBatchA == 0 or (attrs.transA == 0 and ranks.A == 3)", "attrs.transBatchB == 0", "ranks.A >= 2", "ranks.B == ranks.A", "ranks.Y == ranks.A", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) == (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else dim(shapes.B, ranks.B - 2))", "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) >= tunables.SUBGROUP_MATRIX_MIN_M", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) % 32 == 0", "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) % 64 == 0", "(ranks.A == 2 or (ranks.A == 3 and (dim(shapes.A, 1) if attrs.transBatchA != 0 else dim(shapes.A, 0)) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.B, 0)) or (ranks.A == 4 and dim(shapes.A, 0) == dim(shapes.B, 0) and dim(shapes.A, 1) == dim(shapes.B, 1) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.A, 1)))", "dim(shapes.Y, ranks.Y - 2) == (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))", "dim(shapes.Y, ranks.Y - 1) == (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))", "ceil((dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) * (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"],
-      "constants": {
         "usesF16": "dtypes.T == \"f16\"",
         "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
         "outScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
@@ -249,37 +394,32 @@
         {
           "id": "main",
           "name": "FusedMatMul.SubgroupMatrix",
-          "source": { "shader": "fused-matmul-subgroup-matrix.wgsl.jinja", "inputs": { "alpha": "attrs.alpha" } },
-          "bindings": "matrix",
-          "dispatch": {
-            "x": "ceil(constants.N / 64)",
-            "y": "ceil(constants.M / 32)",
-            "z": "numel(shapes.Y) / (constants.M * constants.N)"
-          }
         }
       ]
     },
     {
       "id": "broadcast_rank4_tiled_reg",
       "priority": 6,
-      "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "ranks.A == 4", "(ranks.B == 2 or ranks.B == 3)", "ranks.Y == 4", "dim(shapes.Y, 0) == dim(shapes.A, 0)", "(ranks.B == 2 or dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1)", "dim(shapes.Y, 1) == (dim(shapes.A, 1) if ranks.B == 2 else max(dim(shapes.A, 1), dim(shapes.B, 0)))", "dim(shapes.A, 3) == dim(shapes.B, ranks.B - 2)", "dim(shapes.Y, 2) == dim(shapes.A, 2)", "dim(shapes.Y, 3) == dim(shapes.B, ranks.B - 1)", "dim(shapes.A, 2) >= 64", "dim(shapes.A, 3) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "ceil(dim(shapes.B, ranks.B - 1) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 2) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) <= device.limits.maxComputeWorkgroupsPerDimension"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.BroadcastRank4TiledReg",
-          "source": {
-            "shader": "matmul-tiled-general-reg.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "alpha": "attrs.alpha",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "transBatchA": "false"
-            }
           },
-          "bindings": "matrix",
           "dispatch": {
             "x": "ceil(dim(shapes.B, ranks.B - 1) / 64)",
             "y": "ceil(dim(shapes.A, 2) / 64)",
@@ -290,26 +430,23 @@
     },
     {
       "id": "plain_rank2_tiled_reg",
-      "description": "Register-blocked rank-2 Y = alpha * A @ B for the tier without subgroup matrices. It is the non-transposed rung the two specialised register-blocked variants beside it never covered, so plain rank-2 shapes fell to the scalar-indexed tile.",
       "priority": 4,
-      "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "ranks.A == 2", "ranks.B == 2", "ranks.Y == 2", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 0) == dim(shapes.A, 0)", "dim(shapes.Y, 1) == dim(shapes.B, 1)", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 1) >= 32", "dim(shapes.B, 1) >= 64", "ceil(dim(shapes.A, 0) / 64) * ceil(dim(shapes.B, 1) / 64) >= tunables.TILED_REG_MIN_WORKGROUPS", "ceil(dim(shapes.B, 1) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 0) / 64) <= device.limits.maxComputeWorkgroupsPerDimension"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.PlainRank2TiledReg",
-          "source": {
-            "shader": "matmul-tiled-general-reg.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "alpha": "attrs.alpha",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "transBatchA": "false"
-            }
           },
-          "bindings": "matrix",
           "dispatch": { "x": "ceil(dim(shapes.B, 1) / 64)", "y": "ceil(dim(shapes.A, 0) / 64)", "z": 1 }
         }
       ]
@@ -317,25 +454,23 @@
     {
       "id": "transbatch_a_tiled_reg",
       "priority": 5,
-      "when": ["f16Ok(dtypes.T)", "attrs.transBatchA != 0", "attrs.transBatchB == 0", "attrs.transA == 0", "attrs.transB == 0", "ranks.A == 3", "ranks.B == 3", "ranks.Y == 3", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 0) == dim(shapes.B, 0)", "dim(shapes.A, 2) == dim(shapes.B, 1)", "dim(shapes.Y, 1) == dim(shapes.A, 0)", "dim(shapes.Y, 2) == dim(shapes.B, 2)", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 2) >= 32", "dim(shapes.B, 2) >= 64", "ceil(dim(shapes.B, 2) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 0) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "dim(shapes.Y, 0) <= device.limits.maxComputeWorkgroupsPerDimension"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.TransBatchATiledReg",
-          "source": {
-            "shader": "matmul-tiled-general-reg.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "alpha": "attrs.alpha",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "transBatchA": "true",
-              "kTile": "4"
-            }
           },
-          "bindings": "matrix",
           "dispatch": {
             "x": "ceil(dim(shapes.B, 2) / 64)",
             "y": "ceil(dim(shapes.A, 0) / 64)",
@@ -347,27 +482,25 @@
     {
       "id": "tiled",
       "priority": 0,
-      "when": ["ranks.A >= 1", "ranks.B >= 1", "f16Ok(dtypes.T)", "transBatchContract", "(dim(shapes.A, 0) if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA == 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) == (dim(shapes.B, 0) if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2))))", "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) / 16) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))) / 16) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / max(1, (1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) * (1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2))))) <= device.limits.maxComputeWorkgroupsPerDimension"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.Tiled",
-          "source": {
-            "shader": "matmul-tiled-general.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "transA": "attrs.transA != 0",
-              "transB": "attrs.transB != 0",
-              "alpha": "attrs.alpha",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "transBatchA": "attrs.transBatchA != 0",
-              "transBatchB": "attrs.transBatchB != 0"
-            }
           },
-          "bindings": "matrix",
           "dispatch": {
             "x": "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))) / 32)",
             "y": "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) / 32)",
@@ -375,6 +508,35 @@
           }
         }
       ]
     }
   ]
 }

   "domain": "com.microsoft",
   "name": "FusedMatMul",
   "sinceVersion": 1,
+  "inputs": { "A": { "dtype": "T" }, "B": { "dtype": "T" } },
+  "outputs": {
+    "Y": {
       "dtype": "T",
       "rank": "max(ranks.A, ranks.B) - (1 if ranks.A == 1 or ranks.B == 1 else 0)",
+      "shape": "matmulShape(logicalAShape, logicalBShape)"
     }
   },
+  "attributes": {
+    "alpha": { "default": 1 },
+    "transA": { "default": 0 },
+    "transB": { "default": 0 },
+    "transBatchA": { "default": 0 },
+    "transBatchB": { "default": 0 }
   },
+  "typeConstraints": { "T": ["float32", "float16"] },
   "tunables": {
+    "TILED_REG_MIN_WORKGROUPS": { "default": 64 },
+    "GEMV_TARGET_BLOCKS": { "default": 512 },
+    "SUBGROUP_MATRIX_MIN_M": { "default": 2 },
+    "SUBGROUP_MATRIX_SPLITK_TARGET_WGS": { "default": 512 },
+    "SUBGROUP_MATRIX_SPLITK_MIN_K": { "default": 1024 },
+    "SUBGROUP_MATRIX_SPLITK_MAX_TILES": { "default": 128 },
+    "BAND_VEC4_MAX_ROWS": { "default": 16 },
+    "BAND_SPLIT_TARGET_WORKGROUPS": { "default": 256 },
+    "BAND_SPLIT_MAX_COLUMN_GROUPS": { "default": 24 },
+    "BAND_SPLIT_SLICES": { "default": 8 },
+    "TRANSBATCH_B_SUBGROUP_MATRIX_MIN_K": { "default": 128 },
+    "TRANSBATCH_B_SUBGROUP_MATRIX_MIN_WORKGROUPS": { "default": 64 },
+    "TRANSBATCH_REG_MIN_K": { "default": 128 },
+    "BAND_PREFER_MAX_ROWS": { "default": 8 },
+    "BAND_PREFER_DEEP_K": { "default": 4096 }
   },
   "derive": {
     "gemvWorkgroups": "ceilDiv(dim(shapes.B, 1), 128)",
     "gemvSliceCap": "min(32, device.limits.maxComputeWorkgroupSizeY, floor(device.limits.maxComputeInvocationsPerWorkgroup / 32), floor(device.limits.maxComputeWorkgroupStorageSize / 512))",
+    "gemvSlicesPlan": "max(1, min(gemvSliceCap, max(8, pow2ceil(ceilDiv(tunables.GEMV_TARGET_BLOCKS, gemvWorkgroups)))))",
     "batchMovedAShape": "moveAxis(shapes.A, 0, -2) if attrs.transBatchA != 0 else shapes.A",
     "batchMovedBShape": "moveAxis(shapes.B, 0, -2) if attrs.transBatchB != 0 else shapes.B",
     "logicalAShape": "moveAxis(batchMovedAShape, -1, -2) if attrs.transA != 0 and ranks.A > 1 else batchMovedAShape",
     "sgmatSplitK8Ok": "sgmatSplitKDepth % 256 == 0",
     "sgmatSplitK4Ok": "sgmatSplitKDepth % 128 == 0",
     "sgmatSplitK2Ok": "sgmatSplitKDepth % 64 == 0",
+    "sgmatSplitK": "32 if (sgmatSplitKWant > 16 and sgmatSplitK32Ok) else (16 if (sgmatSplitKWant > 8 and sgmatSplitK16Ok) else (8 if (sgmatSplitKWant > 4 and sgmatSplitK8Ok) else (4 if (sgmatSplitKWant > 2 and sgmatSplitK4Ok) else (2 if sgmatSplitK2Ok else 1))))",
+    "bandRank2Ok": "ranks.A == 2 and ranks.B == 2 and ranks.Y == 2 and attrs.transA == 0 and attrs.transB == 0 and attrs.transBatchA == 0 and attrs.transBatchB == 0 and dim(shapes.A, 1) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.B, 1)",
+    "bandSplitWant": "pow2ceil(ceilDiv(tunables.BAND_SPLIT_TARGET_WORKGROUPS, max(1, gemvWorkgroups)))",
+    "bandSplitK": "16 if (bandSplitWant >= 16 and dim(shapes.A, ranks.A - 1) >= 4096) else (8 if (bandSplitWant >= 8 and dim(shapes.A, ranks.A - 1) >= 2048) else (4 if (bandSplitWant >= 4 and dim(shapes.A, ranks.A - 1) >= 1024) else (2 if (bandSplitWant >= 2 and dim(shapes.A, ranks.A - 1) >= 512) else 1)))"
+  },
+  "bindings": {
+    "a": { "arg": "A", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "b": { "arg": "B", "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "partials": { "buffer": "read-only-storage", "elementType": "f32" },
+    "y": { "arg": "Y", "buffer": "storage", "elementType": "$scalar" },
+    "params": { "buffer": "uniform", "struct": [{ "name": "cols", "type": "u32", "value": "numel(shapes.Y)" }] },
+    "b_3": { "arg": "B", "name": "b", "buffer": "read-only-storage", "elementType": "$scalar" }
   },
   "variants": [
+    {
+      "id": "subgroup_matrix_transbatch_b_f16",
+      "priority": 11,
+      "when": ["dtypes.T == \"f16\" and f16Ok(dtypes.T)", "ranks.A == 3 and ranks.B == 3 and ranks.Y == 3", "attrs.transA == 0 and attrs.transB == 0 and attrs.transBatchA == 0 and attrs.transBatchB != 0", "dim(shapes.A, 0) == dim(shapes.B, 1) and dim(shapes.Y, 0) == dim(shapes.A, 0)", "dim(shapes.A, 2) == dim(shapes.B, 0)", "dim(shapes.Y, 1) == dim(shapes.A, 1) and dim(shapes.Y, 2) == dim(shapes.B, 2)", "dim(shapes.A, 0) > 0", "dim(shapes.A, 1) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, 2) >= 32", "dim(shapes.B, 2) >= 64", "ceilDiv(dim(shapes.B, 2), 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.A, 1), 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dim(shapes.A, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "subgroupMatrixResourcesFit", "wave32Effective", "dim(shapes.A, 1) % 32 == 0", "dim(shapes.A, 2) % 32 == 0", "dim(shapes.B, 2) % 64 == 0", "dim(shapes.A, 2) >= tunables.TRANSBATCH_B_SUBGROUP_MATRIX_MIN_K", "dim(shapes.A, 0) * ceilDiv(dim(shapes.A, 1), 32) * ceilDiv(dim(shapes.B, 2), 64) >= tunables.TRANSBATCH_B_SUBGROUP_MATRIX_MIN_WORKGROUPS"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "hasBias": false,
+        "usesF16": "dtypes.T == \"f16\"",
+        "fScalar": "dtypes.T",
+        "outScalar": "dtypes.T",
+        "scalar": "dtypes.T",
+        "generalAddressing": true,
+        "outputBuffer": "\"y\"",
+        "alpha": "attrs.alpha",
+        "M": "dim(shapes.A, ranks.A - 2)",
+        "K": "dim(shapes.A, ranks.A - 1)",
+        "N": "dim(shapes.B, ranks.B - 1)",
+        "batchCount": "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1))"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.SubgroupMatrixTransBatchB",
+          "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.A",
+            "bShape": "logicalBShape",
+            "aRank": "ranks.A",
+            "bRank": "ranks.B",
+            "bStorageStrides": "[dim(shapes.B, 2), dim(shapes.B, 1) * dim(shapes.B, 2), 1]"
+          },
+          "bindings": ["a", "b_3", "y"],
+          "dispatch": { "x": "ceil(N / 64)", "y": "ceil(M / 32)", "z": "batchCount" }
+        }
+      ]
+    },
+    {
+      "id": "subgroup_matrix_transbatch_b_f32",
+      "priority": 11,
+      "when": ["dtypes.T == \"f32\" and f16Ok(dtypes.T)", "ranks.A == 3 and ranks.B == 3 and ranks.Y == 3", "attrs.transA == 0 and attrs.transB == 0 and attrs.transBatchA == 0 and attrs.transBatchB != 0", "dim(shapes.A, 0) == dim(shapes.B, 1) and dim(shapes.Y, 0) == dim(shapes.A, 0)", "dim(shapes.A, 2) == dim(shapes.B, 0)", "dim(shapes.Y, 1) == dim(shapes.A, 1) and dim(shapes.Y, 2) == dim(shapes.B, 2)", "dim(shapes.A, 0) > 0", "dim(shapes.A, 1) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, 2) >= 32", "dim(shapes.B, 2) >= 64", "ceilDiv(dim(shapes.B, 2), 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.A, 1), 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dim(shapes.A, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "subgroupMatrixResourcesFit", "wave32Effective", "dim(shapes.A, 1) % 32 == 0", "dim(shapes.A, 2) % 32 == 0", "dim(shapes.B, 2) % 64 == 0", "dim(shapes.A, 2) >= tunables.TRANSBATCH_B_SUBGROUP_MATRIX_MIN_K", "dim(shapes.A, 0) * ceilDiv(dim(shapes.A, 1), 32) * ceilDiv(dim(shapes.B, 2), 64) >= tunables.TRANSBATCH_B_SUBGROUP_MATRIX_MIN_WORKGROUPS"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }]
+      },
+      "derive": {
+        "hasBias": false,
+        "usesF16": "dtypes.T == \"f16\"",
+        "fScalar": "dtypes.T",
+        "outScalar": "dtypes.T",
+        "scalar": "dtypes.T",
+        "generalAddressing": true,
+        "outputBuffer": "\"y\"",
+        "alpha": "attrs.alpha",
+        "M": "dim(shapes.A, ranks.A - 2)",
+        "K": "dim(shapes.A, ranks.A - 1)",
+        "N": "dim(shapes.B, ranks.B - 1)",
+        "batchCount": "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1))"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.SubgroupMatrixTransBatchB",
+          "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.A",
+            "bShape": "logicalBShape",
+            "aRank": "ranks.A",
+            "bRank": "ranks.B",
+            "bStorageStrides": "[dim(shapes.B, 2), dim(shapes.B, 1) * dim(shapes.B, 2), 1]"
+          },
+          "bindings": ["a", "b_3", "y"],
+          "dispatch": { "x": "ceil(N / 64)", "y": "ceil(M / 32)", "z": "batchCount" }
+        }
+      ]
+    },
     {
       "id": "f32_m1_gemv_vec4",
       "priority": 30,
+      "when": ["dtypes.T == \"f32\"", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "ranks.A == 2", "ranks.B == 2", "ranks.Y == 2", "dim(shapes.A, 0) == 1", "dim(shapes.Y, 0) == 1", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 1) == dim(shapes.B, 1)", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "ceil(dim(shapes.B, 1) / 128) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": { "usesF16": false, "gemvSlices": "gemvSlicesPlan", "alphaScale": "attrs.alpha" },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.F32M1GemvVec4",
           "shader": "matmul-vector-matrix-vec4.wgsl.jinja",
+          "bindings": [
+            { "arg": "A", "name": "a", "elementType": "f32" },
+            { "arg": "B", "name": "b", "elementType": "vec4<f32>" },
+            { "arg": "Y", "name": "c", "elementType": "vec4<f32>" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "K", "type": "u32", "value": "dim(shapes.A, 1)" },
+                { "name": "N4", "type": "u32", "value": "dim(shapes.B, 1) / 4" }
+              ]
+            }
+          ],
           "dispatch": { "x": "ceil(dim(shapes.B, 1) / 128)" }
         }
       ]
     },
+    {
+      "id": "rank2_band_vec4_splitk",
+      "priority": 11,
+      "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)", "bandRank2Ok", "dim(shapes.A, 0) >= 2", "dim(shapes.A, 0) <= tunables.BAND_VEC4_MAX_ROWS", "dim(shapes.A, 1) > 0", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "32 <= device.limits.maxComputeWorkgroupSizeX", "gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS", "bandSplitK >= 2", "bandSplitK * numel(shapes.Y) * 4 <= device.limits.maxStorageBufferBindingSize", "bandSplitK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tunables.BAND_SPLIT_SLICES <= device.limits.maxComputeWorkgroupSizeY", "32 * tunables.BAND_SPLIT_SLICES <= device.limits.maxComputeInvocationsPerWorkgroup"],
+      "derive": {
+        "usesF16": "dtypes.T == \"f16\"",
+        "scalar": "dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "batched": false,
+        "outputBuffer": "\"y\"",
+        "alpha": "attrs.alpha",
+        "M": "dim(shapes.A, 0)",
+        "K": "dim(shapes.A, 1)",
+        "N": "dim(shapes.B, 1)",
+        "gemvSlices": "tunables.BAND_SPLIT_SLICES",
+        "kSplits": "bandSplitK",
+        "split": "bandSplitK",
+        "workgroupSize": 256
+      },
+      "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[bandSplitK * numel(shapes.Y)]" }],
+      "passes": [
+        {
+          "id": "partial",
+          "name": "FusedMatMul.Rank2BandVec4SplitK",
+          "shader": "matmul-band-vec4.wgsl.jinja",
+          "bindings": ["a", "b", { "scratch": "partials", "name": "y", "elementType": "vec4<f32>" }],
+          "dispatch": { "x": "gemvWorkgroups", "y": "bandSplitK" }
+        },
+        {
+          "id": "combine",
+          "name": "FusedMatMul.Rank2BandVec4SplitKCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"", "intMode": false },
+          "bindings": ["partials", "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.Y)), (256)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.Y)), (256)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "rank2_band_vec4",
+      "priority": 11,
+      "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)", "bandRank2Ok", "dim(shapes.A, 0) >= 2", "dim(shapes.A, 0) <= tunables.BAND_VEC4_MAX_ROWS", "dim(shapes.A, 1) > 0", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "32 <= device.limits.maxComputeWorkgroupSizeX", "gemvSlicesPlan <= device.limits.maxComputeWorkgroupSizeY", "32 * gemvSlicesPlan <= device.limits.maxComputeInvocationsPerWorkgroup", "32 * gemvSlicesPlan * 16 <= device.limits.maxComputeWorkgroupStorageSize", "not (gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS and bandSplitK >= 2)"],
+      "derive": {
+        "usesF16": "dtypes.T == \"f16\"",
+        "scalar": "dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "batched": false,
+        "outputBuffer": "\"y\"",
+        "alpha": "attrs.alpha",
+        "M": "dim(shapes.A, 0)",
+        "K": "dim(shapes.A, 1)",
+        "N": "dim(shapes.B, 1)",
+        "gemvSlices": "gemvSlicesPlan"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.Rank2BandVec4",
+          "shader": "matmul-band-vec4.wgsl.jinja",
+          "bindings": ["a", "b", { "arg": "Y", "name": "y", "elementType": "$vectorScalar" }],
+          "dispatch": { "x": "gemvWorkgroups" }
+        }
+      ],
+      "demoteWhen": ["false"]
+    },
+    {
+      "id": "rank2_band_vec4_f32_preferred",
+      "priority": 13,
+      "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)", "bandRank2Ok", "dim(shapes.A, 0) >= 2", "dim(shapes.A, 0) <= tunables.BAND_VEC4_MAX_ROWS", "dim(shapes.A, 1) > 0", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "32 <= device.limits.maxComputeWorkgroupSizeX", "gemvSlicesPlan <= device.limits.maxComputeWorkgroupSizeY", "32 * gemvSlicesPlan <= device.limits.maxComputeInvocationsPerWorkgroup", "32 * gemvSlicesPlan * 16 <= device.limits.maxComputeWorkgroupStorageSize", "not (gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS and bandSplitK >= 2)"],
+      "derive": {
+        "usesF16": "dtypes.T == \"f16\"",
+        "scalar": "dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "batched": false,
+        "outputBuffer": "\"y\"",
+        "alpha": "attrs.alpha",
+        "M": "dim(shapes.A, 0)",
+        "K": "dim(shapes.A, 1)",
+        "N": "dim(shapes.B, 1)",
+        "gemvSlices": "gemvSlicesPlan"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.Rank2BandVec4",
+          "shader": "matmul-band-vec4.wgsl.jinja",
+          "bindings": ["a", "b", { "arg": "Y", "name": "y", "elementType": "$vectorScalar" }],
+          "dispatch": { "x": "gemvWorkgroups" }
+        }
+      ],
+      "demoteWhen": ["dtypes.T != \"f32\" or (dim(shapes.A, 0) > tunables.BAND_PREFER_MAX_ROWS and dim(shapes.A, 1) >= tunables.BAND_PREFER_DEEP_K)"]
+    },
     {
       "id": "subgroup_matrix_splitk",
       "priority": 12,
+      "when": ["(dtypes.T == \"f16\" or dtypes.T == \"f32\") and f16Ok(dtypes.T)", "fusedSgmatRank2Ok", "dim(shapes.A, 0) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, 1) >= tunables.SUBGROUP_MATRIX_SPLITK_MIN_K", "dim(shapes.B, 1) % 64 == 0", "sgmatSplitK >= 2", "sgmatOutTiles < tunables.SUBGROUP_MATRIX_SPLITK_MAX_TILES", "sgmatSplitK * numel(shapes.Y) * 4 <= device.limits.maxStorageBufferBindingSize", "sgmatSplitK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.Y, 1), 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.Y, 0), 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "subgroupMatrixResourcesFit", "wave32Effective"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [
           { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
         ]
       },
+      "derive": {
         "usesF16": "dtypes.T == \"f16\"",
         "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
         "scalar": "dtypes.T",
         {
           "id": "partial",
           "name": "FusedMatMul.SubgroupMatrixSplitK",
+          "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": {
+            "aShape": ["dim(shapes.A, 0)", "dim(shapes.A, 1)"],
+            "bShape": ["dim(shapes.B, 0)", "dim(shapes.B, 1)"],
+            "aRank": 2,
+            "bRank": 2
           },
+          "bindings": ["a", "b_3", { "name": "partials", "buffer": "storage", "elementType": "f32" }],
           "dispatch": { "x": "ceilDiv(dim(shapes.Y, 1), 64)", "y": "ceilDiv(dim(shapes.Y, 0), 32)", "z": "sgmatSplitK" }
         },
         {
           "id": "combine",
           "name": "FusedMatMul.SubgroupMatrixSplitKCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"", "intMode": false },
+          "bindings": ["partials", "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.Y)), (256)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.Y)), (256)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "subgroup_matrix_tail_broadcast",
       "priority": 11,
+      "when": ["dtypes.T == \"f16\"", "f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "(((ranks.A == 2 and ranks.B == 2 and ranks.Y == 2) or (ranks.A == 4 and ranks.B == 3 and ranks.Y == 4 and dim(shapes.Y, 0) == dim(shapes.A, 0) and (dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1) and dim(shapes.Y, 1) == max(dim(shapes.A, 1), dim(shapes.B, 0)))) or (ranks.A == 4 and ranks.B == 2 and ranks.Y == 4 and sameShape(prefix(shapes.Y, 2), prefix(shapes.A, 2))))", "dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 2)", "dim(shapes.A, ranks.A - 2) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, ranks.A - 1) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "dim(shapes.Y, ranks.Y - 2) == dim(shapes.A, ranks.A - 2)", "dim(shapes.Y, ranks.Y - 1) == dim(shapes.B, ranks.B - 1)", "ceil(dim(shapes.B, ranks.B - 1) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, ranks.A - 2) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "wave32Effective"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
       },
+      "derive": {
         "hasBias": false,
         "usesF16": true,
         "fScalar": "\"f16\"",
         {
           "id": "main",
           "name": "FusedMatMul.SubgroupMatrixTailBroadcast",
+          "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": { "aShape": "shapes.A", "bShape": "shapes.B", "aRank": "ranks.A", "bRank": "ranks.B" },
+          "bindings": ["a", "b_3", "y"],
+          "dispatch": { "x": "ceil(N / 64)", "y": "ceil(M / 32)", "z": "batchCount" }
         }
       ]
     },
     {
       "id": "subgroup_matrix",
       "priority": 10,
+      "when": ["f16Ok(dtypes.T)", "attrs.transBatchA == 0 or (attrs.transA == 0 and ranks.A == 3)", "attrs.transBatchB == 0", "ranks.A >= 2", "ranks.B == ranks.A", "ranks.Y == ranks.A", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) == (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else dim(shapes.B, ranks.B - 2))", "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) >= tunables.SUBGROUP_MATRIX_MIN_M", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) % 32 == 0", "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) % 64 == 0", "(ranks.A == 2 or (ranks.A == 3 and (dim(shapes.A, 1) if attrs.transBatchA != 0 else dim(shapes.A, 0)) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.B, 0)) or (ranks.A == 4 and dim(shapes.A, 0) == dim(shapes.B, 0) and dim(shapes.A, 1) == dim(shapes.B, 1) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.A, 1)))", "dim(shapes.Y, ranks.Y - 2) == (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))", "dim(shapes.Y, ranks.Y - 1) == (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))", "ceil((dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) * (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "wave32Effective"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [
           { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
         ]
       },
+      "derive": {
         "usesF16": "dtypes.T == \"f16\"",
         "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
         "outScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
         {
           "id": "main",
           "name": "FusedMatMul.SubgroupMatrix",
+          "shader": "fused-matmul-subgroup-matrix.wgsl.jinja",
+          "derive": { "alpha": "attrs.alpha" },
+          "bindings": ["a", "b_3", "y"],
+          "dispatch": { "x": "ceil(N / 64)", "y": "ceil(M / 32)", "z": "numel(shapes.Y) / (M * N)" }
         }
       ]
     },
     {
       "id": "broadcast_rank4_tiled_reg",
       "priority": 6,
+      "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "ranks.A == 4", "(ranks.B == 2 or ranks.B == 3)", "ranks.Y == 4", "dim(shapes.Y, 0) == dim(shapes.A, 0)", "(ranks.B == 2 or dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1)", "dim(shapes.Y, 1) == (dim(shapes.A, 1) if ranks.B == 2 else max(dim(shapes.A, 1), dim(shapes.B, 0)))", "dim(shapes.A, 3) == dim(shapes.B, ranks.B - 2)", "dim(shapes.Y, 2) == dim(shapes.A, 2)", "dim(shapes.Y, 3) == dim(shapes.B, ranks.B - 1)", "dim(shapes.A, 2) >= 64", "dim(shapes.A, 3) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "ceil(dim(shapes.B, ranks.B - 1) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, 2) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.BroadcastRank4TiledReg",
+          "shader": "matmul-tiled-general-reg.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.A",
+            "bShape": "shapes.B",
+            "alpha": "attrs.alpha",
+            "aRank": "ranks.A",
+            "bRank": "ranks.B",
+            "transBatchA": "false"
           },
+          "bindings": ["a", "b_3", "y"],
           "dispatch": {
             "x": "ceil(dim(shapes.B, ranks.B - 1) / 64)",
             "y": "ceil(dim(shapes.A, 2) / 64)",
     },
     {
       "id": "plain_rank2_tiled_reg",
       "priority": 4,
+      "when": ["f16Ok(dtypes.T)", "fusedSgmatRank2Ok", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 1) >= 32", "dim(shapes.B, 1) >= 64", "ceil(dim(shapes.A, 0) / 64) * ceil(dim(shapes.B, 1) / 64) >= tunables.TILED_REG_MIN_WORKGROUPS", "ceil(dim(shapes.B, 1) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, 0) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.PlainRank2TiledReg",
+          "shader": "matmul-tiled-general-reg.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.A",
+            "bShape": "shapes.B",
+            "alpha": "attrs.alpha",
+            "aRank": "ranks.A",
+            "bRank": "ranks.B",
+            "transBatchA": "false"
           },
+          "bindings": ["a", "b_3", "y"],
           "dispatch": { "x": "ceil(dim(shapes.B, 1) / 64)", "y": "ceil(dim(shapes.A, 0) / 64)", "z": 1 }
         }
       ]
     {
       "id": "transbatch_a_tiled_reg",
       "priority": 5,
+      "when": ["f16Ok(dtypes.T)", "attrs.transBatchA != 0", "attrs.transBatchB == 0", "attrs.transA == 0", "attrs.transB == 0", "ranks.A == 3", "ranks.B == 3", "ranks.Y == 3", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 0) == dim(shapes.B, 0)", "dim(shapes.A, 2) == dim(shapes.B, 1)", "dim(shapes.Y, 1) == dim(shapes.A, 0)", "dim(shapes.Y, 2) == dim(shapes.B, 2)", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 2) >= 32", "dim(shapes.B, 2) >= 64", "ceil(dim(shapes.B, 2) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, 0) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dim(shapes.Y, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.TransBatchATiledReg",
+          "shader": "matmul-tiled-general-reg.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.A",
+            "bShape": "shapes.B",
+            "alpha": "attrs.alpha",
+            "aRank": "ranks.A",
+            "bRank": "ranks.B",
+            "transBatchA": "true",
+            "kTile": "4"
           },
+          "bindings": ["a", "b_3", "y"],
           "dispatch": {
             "x": "ceil(dim(shapes.B, 2) / 64)",
             "y": "ceil(dim(shapes.A, 0) / 64)",
     {
       "id": "tiled",
       "priority": 0,
+      "when": ["ranks.A >= 1", "ranks.B >= 1", "f16Ok(dtypes.T)", "transBatchContract", "(dim(shapes.A, 0) if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA == 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) == (dim(shapes.B, 0) if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2))))", "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) / 16) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))) / 16) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / max(1, (1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) * (1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2))))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.Tiled",
+          "shader": "matmul-tiled-general.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.A",
+            "bShape": "shapes.B",
+            "transA": "attrs.transA != 0",
+            "transB": "attrs.transB != 0",
+            "alpha": "attrs.alpha",
+            "aRank": "ranks.A",
+            "bRank": "ranks.B",
+            "transBatchA": "attrs.transBatchA != 0",
+            "transBatchB": "attrs.transBatchB != 0"
           },
+          "bindings": ["a", "b_3", "y"],
           "dispatch": {
             "x": "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))) / 32)",
             "y": "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) / 32)",
           }
         }
       ]
+    },
+    {
+      "id": "transbatch_b_tiled_reg",
+      "priority": 5,
+      "when": ["f16Ok(dtypes.T)", "ranks.A == 3 and ranks.B == 3 and ranks.Y == 3", "attrs.transA == 0 and attrs.transB == 0 and attrs.transBatchA == 0 and attrs.transBatchB != 0", "dim(shapes.A, 0) == dim(shapes.B, 1) and dim(shapes.Y, 0) == dim(shapes.A, 0)", "dim(shapes.A, 2) == dim(shapes.B, 0)", "dim(shapes.Y, 1) == dim(shapes.A, 1) and dim(shapes.Y, 2) == dim(shapes.B, 2)", "dim(shapes.A, 0) > 0", "dim(shapes.A, 1) >= 64", "dim(shapes.A, 2) >= tunables.TRANSBATCH_REG_MIN_K", "dim(shapes.B, 2) >= 64", "ceilDiv(dim(shapes.B, 2), 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.A, 1), 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dim(shapes.A, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dim(shapes.A, 0) * ceilDiv(dim(shapes.A, 1), 64) * ceilDiv(dim(shapes.B, 2), 64) >= tunables.TILED_REG_MIN_WORKGROUPS", "16 <= device.limits.maxComputeWorkgroupSizeX and 16 <= device.limits.maxComputeWorkgroupSizeY and 256 <= device.limits.maxComputeInvocationsPerWorkgroup", "(64 * 16 + 64 * 16) * dtypeBytes(dtypes.T) <= device.limits.maxComputeWorkgroupStorageSize"],
+      "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.TransBatchBTiledReg",
+          "shader": "matmul-tiled-general-reg.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.A",
+            "bShape": "logicalBShape",
+            "alpha": "attrs.alpha",
+            "aRank": "ranks.A",
+            "bRank": "ranks.B",
+            "transBatchA": false,
+            "bStorageStrides": "[dim(shapes.B, 2), dim(shapes.B, 1) * dim(shapes.B, 2), 1]",
+            "regSequentialK": "dtypes.T == \"f16\""
+          },
+          "bindings": ["a", "b_3", "y"],
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.B, 2), 64)",
+            "y": "ceilDiv(dim(shapes.A, 1), 64)",
+            "z": "dim(shapes.Y, 0)"
+          }
+        }
+      ]
     }
   ]
 }

build/webgpu/matmul-band-vec4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,84 @@

+// Band GEMM y[M, N] = a[M, K] @ B[K, N]. Each lane owns one vec4 column group
+// and carries one accumulator per row, so a loaded B word feeds all M row
+// accumulators and each A value is reused across four adjacent output columns.
+//
+// A batched consumer runs one band per workgroup row: workgroup_id.y selects
+// the matrix, and every operand is offset by its per-matrix extent.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const K: u32 = {{ K }}u;
+const N4: u32 = {{ N }}u / 4u;
+const LANES: u32 = 32u;
+// SLICES partitions the K reduction across the workgroup's second dimension.
+const SLICES: u32 = {{ gemvSlices }}u;
+{% set kSplitsValue = kSplits if kSplits is defined else 1 %}
+{% set alphaValue = alpha if alpha is defined else 1 %}
+{% if kSplitsValue > 1 %}
+const K_SPLITS: u32 = {{ kSplitsValue }}u;
+const K_PER_SPLIT: u32 = (K + K_SPLITS - 1u) / K_SPLITS;
+{% endif %}
+// The rows drain through one 32 x SLICES array in turn, so the workgroup
+// footprint does not grow with the band.
+var<workgroup> partials: array<vec4<f32>, LANES * SLICES>;
+@compute @workgroup_size(32, {{ gemvSlices }}, 1)
+fn main(
+  @builtin(workgroup_id) workgroup_id: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let lane = lid.x;
+  let slice = lid.y;
+  let cg = workgroup_id.x * LANES + lane;
+{% if kSplitsValue > 1 %}
+  let a_base = 0u;
+  let b_base = 0u;
+  let y_base = workgroup_id.y * ({{ M }}u * N4);
+  let k_begin = workgroup_id.y * K_PER_SPLIT;
+  let k_end = min(K, k_begin + K_PER_SPLIT);
+{% else %}
+  let a_base = 0u;
+  let b_base = 0u;
+  let y_base = 0u;
+{% endif %}
+{% for r in range(M) %}
+  var acc{{ r }} = vec4<f32>(0.0);
+{% endfor %}
+  if (cg < N4) {
+{% if kSplitsValue > 1 %}
+    for (var k = k_begin + slice; k < k_end; k = k + SLICES) {
+{% else %}
+    for (var k = slice; k < K; k = k + SLICES) {
+{% endif %}
+      let bv = vec4<f32>(b[b_base + k * N4 + cg]);
+{% for r in range(M) %}
+      acc{{ r }} = acc{{ r }} + f32(a[a_base + {{ r }}u * K + k]) * bv;
+{% endfor %}
+    }
+  }
+{% for r in range(M) %}
+  partials[slice * LANES + lane] = acc{{ r }};
+  workgroupBarrier();
+  if (slice == 0u && cg < N4) {
+    var total = partials[lane];
+    for (var s = 1u; s < SLICES; s = s + 1u) {
+      total = total + partials[s * LANES + lane];
+    }
+{% if alphaValue != 1 %}
+    total = total * {{ alphaValue }};
+{% endif %}
+{% if kSplitsValue > 1 %}
+    {{ outputBuffer }}[y_base + {{ r }}u * N4 + cg] = total;
+{% else %}
+    {{ outputBuffer }}[y_base + {{ r }}u * N4 + cg] = vec4<{{ T }}>(total);
+{% endif %}
+  }
+{% if not loop.last %}
+  workgroupBarrier();
+{% endif %}
+{% endfor %}
+}

build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja CHANGED Viewed

@@ -1,9 +1,6 @@
 // Subgroup-matrix matmul over row-major, batch-outermost operands, with alpha,
 // dense/broadcast batching and guarded K/N tails under `generalAddressing`, and
 // an optional fused bias on the direct dense path that omits it.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
@@ -11,6 +8,7 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = fScalar %}
@@ -20,17 +18,16 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 {% set SPLIT_K = splitK if splitK is defined else 1 %}
 {% set OUT = outputBuffer if outputBuffer is defined else "c" %}
 {% set OUT_SCALAR = outScalar if outScalar is defined else T %}
-{% if GENERAL %}
-{% set aR = source.aRank %}
-{% set bR = source.bRank %}
 {% set aBatchLen = aR - 2 %}
 {% set bBatchLen = bR - 2 %}
 {% set batchRank = aBatchLen %}
-{% set aMStride = source.aShape[aR-1] %}
 {% set aKStride = 1 %}
-{% set bKStride = source.bShape[bR-1] %}
 {% set bNStride = 1 %}
-{% endif %}
 const M: u32 = {{ M }}u;
 const K: u32 = {{ K }}u;
@@ -40,7 +37,6 @@ const BATCH_COUNT: u32 = {{ batchCount if batchCount is defined else 1 }}u;
 const SPLIT_K: u32 = {{ SPLIT_K }}u;
 const K_PER_SPLIT: u32 = {{ kPerSplit }}u;
 {% endif %}
-{% if GENERAL %}
 const A_M_STRIDE: u32 = {{ aMStride }}u;
 const A_K_STRIDE: u32 = {{ aKStride }}u;
 const B_K_STRIDE: u32 = {{ bKStride }}u;
@@ -48,10 +44,6 @@ const B_N_STRIDE: u32 = {{ bNStride }}u;
 {% if TAIL %}const K_FULL: u32 = (K / 32u) * 32u;
 {% endif %}
 const ALPHA: f32 = f32({{ alpha }});
-{% else %}
-const A_BATCH_STRIDE: u32 = M * K;
-const B_BATCH_STRIDE: u32 = K * N;
-{% endif %}
 const C_BATCH_STRIDE: u32 = M * N;
 const TILE_COLS: u32 = 64u;
 const TILE_ROWS: u32 = 32u;
@@ -70,17 +62,9 @@ fn loadSHMA(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
     let k = k_idx + col + col_offset;
     if (a_global < M) {
 {% if operandScalar == "f16" %}
-{% if GENERAL %}
       tile_A[row * TILE_K + col + col_offset] = f16(a[a_base + a_global * A_M_STRIDE + k * A_K_STRIDE]);
 {% else %}
-      tile_A[row * TILE_K + col + col_offset] = f16(a[a_base + a_global * K + k]);
-{% endif %}
-{% else %}
-{% if GENERAL %}
       tile_A[row * TILE_K + col + col_offset] = f32(a[a_base + a_global * A_M_STRIDE + k * A_K_STRIDE]);
-{% else %}
-      tile_A[row * TILE_K + col + col_offset] = f32(a[a_base + a_global * K + k]);
-{% endif %}
 {% endif %}
     } else {
 {% if operandScalar == "f16" %}
@@ -112,7 +96,6 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   let col = c_idx * 16u;
   for (var i = 0u; i < 16u; i = i + 1u) {
     let k = k_idx + col + i;
-{% if GENERAL %}
 {% if TAIL %}
     let b_safe = min(b_col, N - 1u);
     tile_B[row * TILE_K + col + i] = {{ operandScalar }}(b[b_base + k * B_K_STRIDE + b_safe * B_N_STRIDE]);
@@ -122,10 +105,6 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
 {% else %}
     tile_B[row * TILE_K + col + i] = f32(b[b_base + k * B_K_STRIDE + b_col * B_N_STRIDE]);
 {% endif %}
-{% endif %}
-{% else %}
-{% set bIndex = "b_col * K + k" if bTransposed is defined and bTransposed else "k * N + b_col" %}
-    tile_B[row * TILE_K + col + i] = {{ operandScalar }}(b[b_base + {{ bIndex }}]);
 {% endif %}
   }
 }
@@ -149,7 +128,6 @@ fn loadSHMBKTail(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32)
 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 = col + 1u;
-{% if GENERAL %}
 {% for block in range(4) %}
 {% if TAIL %}
     if (col_base + col + {{ block * 8 }}u < N) {
@@ -168,29 +146,6 @@ fn storeOutput(offset: u32{% if needsColBase %}, col_base: u32{% endif %}, row:
     }
 {% endif %}
 {% endfor %}
-{% else %}
-{% if hasBias %}
-    {{ OUT }}[offset + row * N + col]       = {{ T }}(scratch[src_slot][0][row * 8u + col] + {{ accScalar }}(bias[col_base + col]));
-    {{ OUT }}[offset + row * N + col + 8u]  = {{ T }}(scratch[src_slot][1][row * 8u + col] + {{ accScalar }}(bias[col_base + col + 8u]));
-    {{ OUT }}[offset + row * N + col + 16u] = {{ T }}(scratch[src_slot][2][row * 8u + col] + {{ accScalar }}(bias[col_base + col + 16u]));
-    {{ OUT }}[offset + row * N + col + 24u] = {{ T }}(scratch[src_slot][3][row * 8u + col] + {{ accScalar }}(bias[col_base + col + 24u]));
-    {{ OUT }}[offset + row * N + col2]       = {{ T }}(scratch[src_slot][0][row * 8u + col2] + {{ accScalar }}(bias[col_base + col2]));
-    {{ OUT }}[offset + row * N + col2 + 8u]  = {{ T }}(scratch[src_slot][1][row * 8u + col2] + {{ accScalar }}(bias[col_base + col2 + 8u]));
-    {{ OUT }}[offset + row * N + col2 + 16u] = {{ T }}(scratch[src_slot][2][row * 8u + col2] + {{ accScalar }}(bias[col_base + col2 + 16u]));
-    {{ OUT }}[offset + row * N + col2 + 24u] = {{ T }}(scratch[src_slot][3][row * 8u + col2] + {{ accScalar }}(bias[col_base + col2 + 24u]));
-{% else %}
-    {{ OUT }}[offset + row * N + col]       = {{ T }}(scratch[src_slot][0][row * 8u + col]);
-    {{ OUT }}[offset + row * N + col + 8u]  = {{ T }}(scratch[src_slot][1][row * 8u + col]);
-    {{ OUT }}[offset + row * N + col + 16u] = {{ T }}(scratch[src_slot][2][row * 8u + col]);
-    {{ OUT }}[offset + row * N + col + 24u] = {{ T }}(scratch[src_slot][3][row * 8u + col]);
-    {{ OUT }}[offset + row * N + col2]       = {{ T }}(scratch[src_slot][0][row * 8u + col2]);
-    {{ OUT }}[offset + row * N + col2 + 8u]  = {{ T }}(scratch[src_slot][1][row * 8u + col2]);
-    {{ OUT }}[offset + row * N + col2 + 16u] = {{ T }}(scratch[src_slot][2][row * 8u + col2]);
-    {{ OUT }}[offset + row * N + col2 + 24u] = {{ T }}(scratch[src_slot][3][row * 8u + col2]);
-{% endif %}
-{% endif %}
   }
 }
@@ -222,8 +177,8 @@ fn main(
   let row_tile_stride = num_wg.y * TILE_ROWS;
   for (var a_global_base = workgroup_id.y * TILE_ROWS; a_global_base < M; a_global_base += row_tile_stride) {
 {% if SPLIT_K > 1 %}
-  // Split-K maps z to (batch, K segment), multiplying occupancy for narrow
-  // matrices while every segment still owns a disjoint contiguous K range.
   for (var batch_split = workgroup_id.z; batch_split < BATCH_COUNT * SPLIT_K; batch_split += num_wg.z) {
     let batch = batch_split / SPLIT_K;
     let split_id = batch_split - batch * SPLIT_K;
@@ -231,14 +186,13 @@ fn main(
   // workgroup_size.z = 1, so num_wg.z is the dispatch stride over the batch axis.
   for (var batch = workgroup_id.z; batch < BATCH_COUNT; batch += num_wg.z) {
 {% endif %}
-{% if GENERAL %}
 {% set hasBatchCoord = namespace(value=false) %}
 {% for i in range(batchRank) %}
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
-{% set aDim = source.aShape[aAxis] %}
-{% set bDim = source.bShape[bAxis] if bAxis >= 0 else 1 %}
 {% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
 {% endfor %}
     // Right-aligned broadcast offsets, decomposed from the flattened output batch.
@@ -251,13 +205,14 @@ fn main(
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
-{% set aDim = source.aShape[aAxis] %}
-{% set bDim = source.bShape[bAxis] if bAxis >= 0 else 1 %}
 {% set outDim = aDim if aDim >= bDim else bDim %}
 {% set aStride = namespace(v=1) %}
-{% if aDim != 1 %}{% for j in range(aAxis + 1, aR) %}{% set aStride.v = aStride.v * source.aShape[j] %}{% endfor %}{% else %}{% set aStride.v = 0 %}{% endif %}
 {% set bStride = namespace(v=1) %}
-{% if bAxis >= 0 and bDim != 1 %}{% for j in range(bAxis + 1, bR) %}{% set bStride.v = bStride.v * source.bShape[j] %}{% endfor %}{% else %}{% set bStride.v = 0 %}{% endif %}
 {% if outDim > 1 %}
     let c{{ axis }} = zTmp % {{ outDim }}u;
     zTmp = zTmp / {{ outDim }}u;
@@ -272,11 +227,6 @@ fn main(
 {% else %}
     let c_base = batch * C_BATCH_STRIDE;
 {% endif %}
-{% else %}
-    let a_base = batch * A_BATCH_STRIDE;
-    let b_base = batch * B_BATCH_STRIDE;
-    let c_base = batch * C_BATCH_STRIDE;
-{% endif %}
     var matC00: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
     var matC01: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
@@ -299,15 +249,16 @@ fn main(
       workgroupBarrier();
       for (var step = 0u; step < TILE_K; step = step + 8u) {
-        let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
-      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
-      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
-      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
-      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
-      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
-      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -327,15 +278,16 @@ fn main(
       workgroupBarrier();
       for (var step = 0u; step < TILE_K; step = step + 8u) {
-        let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
-      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
-      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
-      let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
-      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
-      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
-      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
-      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);

 // Subgroup-matrix matmul over row-major, batch-outermost operands, with alpha,
 // dense/broadcast batching and guarded K/N tails under `generalAddressing`, and
 // an optional fused bias on the direct dense path that omits it.
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = fScalar %}
 {% set SPLIT_K = splitK if splitK is defined else 1 %}
 {% set OUT = outputBuffer if outputBuffer is defined else "c" %}
 {% set OUT_SCALAR = outScalar if outScalar is defined else T %}
+{% set aR = aRank %}
+{% set bR = bRank %}
 {% set aBatchLen = aR - 2 %}
 {% set bBatchLen = bR - 2 %}
 {% set batchRank = aBatchLen %}
+{% set aMStride = aShape[aR-1] %}
 {% set aKStride = 1 %}
+{% set bKStride = bShape[bR-1] %}
 {% set bNStride = 1 %}
+{% if bStorageStrides is defined %}{% set bKStride = bStorageStrides[bR-2] %}{% set bNStride = bStorageStrides[bR-1] %}{% endif %}
 const M: u32 = {{ M }}u;
 const K: u32 = {{ K }}u;
 const SPLIT_K: u32 = {{ SPLIT_K }}u;
 const K_PER_SPLIT: u32 = {{ kPerSplit }}u;
 {% endif %}
 const A_M_STRIDE: u32 = {{ aMStride }}u;
 const A_K_STRIDE: u32 = {{ aKStride }}u;
 const B_K_STRIDE: u32 = {{ bKStride }}u;
 {% if TAIL %}const K_FULL: u32 = (K / 32u) * 32u;
 {% endif %}
 const ALPHA: f32 = f32({{ alpha }});
 const C_BATCH_STRIDE: u32 = M * N;
 const TILE_COLS: u32 = 64u;
 const TILE_ROWS: u32 = 32u;
     let k = k_idx + col + col_offset;
     if (a_global < M) {
 {% if operandScalar == "f16" %}
       tile_A[row * TILE_K + col + col_offset] = f16(a[a_base + a_global * A_M_STRIDE + k * A_K_STRIDE]);
 {% else %}
       tile_A[row * TILE_K + col + col_offset] = f32(a[a_base + a_global * A_M_STRIDE + k * A_K_STRIDE]);
 {% endif %}
     } else {
 {% if operandScalar == "f16" %}
   let col = c_idx * 16u;
   for (var i = 0u; i < 16u; i = i + 1u) {
     let k = k_idx + col + i;
 {% if TAIL %}
     let b_safe = min(b_col, N - 1u);
     tile_B[row * TILE_K + col + i] = {{ operandScalar }}(b[b_base + k * B_K_STRIDE + b_safe * B_N_STRIDE]);
 {% else %}
     tile_B[row * TILE_K + col + i] = f32(b[b_base + k * B_K_STRIDE + b_col * B_N_STRIDE]);
 {% endif %}
 {% endif %}
   }
 }
 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 = col + 1u;
 {% for block in range(4) %}
 {% if TAIL %}
     if (col_base + col + {{ block * 8 }}u < N) {
     }
 {% endif %}
 {% endfor %}
   }
 }
   let row_tile_stride = num_wg.y * TILE_ROWS;
   for (var a_global_base = workgroup_id.y * TILE_ROWS; a_global_base < M; a_global_base += row_tile_stride) {
 {% if SPLIT_K > 1 %}
+  // Split-K maps z to (batch, K segment), increasing the independent workgroup
+  // count for narrow matrices. Every segment owns a disjoint contiguous K range.
   for (var batch_split = workgroup_id.z; batch_split < BATCH_COUNT * SPLIT_K; batch_split += num_wg.z) {
     let batch = batch_split / SPLIT_K;
     let split_id = batch_split - batch * SPLIT_K;
   // workgroup_size.z = 1, so num_wg.z is the dispatch stride over the batch axis.
   for (var batch = workgroup_id.z; batch < BATCH_COUNT; batch += num_wg.z) {
 {% endif %}
 {% set hasBatchCoord = namespace(value=false) %}
 {% for i in range(batchRank) %}
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aDim = aShape[aAxis] %}
+{% set bDim = bShape[bAxis] if bAxis >= 0 else 1 %}
 {% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
 {% endfor %}
     // Right-aligned broadcast offsets, decomposed from the flattened output batch.
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aDim = aShape[aAxis] %}
+{% set bDim = bShape[bAxis] if bAxis >= 0 else 1 %}
 {% set outDim = aDim if aDim >= bDim else bDim %}
 {% set aStride = namespace(v=1) %}
+{% if aDim != 1 %}{% for j in range(aAxis + 1, aR) %}{% set aStride.v = aStride.v * aShape[j] %}{% endfor %}{% else %}{% set aStride.v = 0 %}{% endif %}
 {% set bStride = namespace(v=1) %}
+{% if bAxis >= 0 and bDim != 1 %}{% for j in range(bAxis + 1, bR) %}{% set bStride.v = bStride.v * bShape[j] %}{% endfor %}{% else %}{% set bStride.v = 0 %}{% endif %}
+{% if bStorageStrides is defined and bAxis >= 0 and bDim != 1 %}{% set bStride.v = bStorageStrides[bAxis] %}{% endif %}
 {% if outDim > 1 %}
     let c{{ axis }} = zTmp % {{ outDim }}u;
     zTmp = zTmp / {{ outDim }}u;
 {% else %}
     let c_base = batch * C_BATCH_STRIDE;
 {% endif %}
     var matC00: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
     var matC01: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
       workgroupBarrier();
       for (var step = 0u; step < TILE_K; step = step + 8u) {
+  {% set directInputs = directMatrixInputs is defined and directMatrixInputs %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+{% for r in range(2) %}
+      var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});
+{% endfor %}
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+{% for c in range(4) %}
+      var matB{{ c }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&{{ "xm" if directInputs else "tile_B" }}, matrix_b_offset{% if c > 0 %} + {{ c * 8 }}u * TILE_K{% endif %}, {{ "N" if directInputs else "TILE_K" }});
+{% endfor %}
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
       workgroupBarrier();
       for (var step = 0u; step < TILE_K; step = step + 8u) {
+  {% set directInputs = directMatrixInputs is defined and directMatrixInputs %}
+      let matrix_a_offset = {% if directInputs %}(a_global_base + subtile_idy * SUB_ROWS) * K + kidx + step{% else %}subtile_idy * SUB_ROWS * TILE_K + step{% endif %};
+{% for r in range(2) %}
+      var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});
+{% endfor %}
+      let matrix_b_offset = {% if directInputs %}b_base + (kidx + step) * N + b_global_base + subtile_idx * SUB_COLS{% else %}subtile_idx * SUB_COLS * TILE_K + step{% endif %};
+{% for c in range(4) %}
+      var matB{{ c }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, {% if directInputs %}row_major{% else %}col_major{% endif %}>(&{{ "xm" if directInputs else "tile_B" }}, matrix_b_offset{% if c > 0 %} + {{ c * 8 }}u{% if not directInputs %} * TILE_K{% endif %}{% endif %}, {{ "N" if directInputs else "TILE_K" }});
+{% endfor %}
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);

build/webgpu/matmul-tiled-general-reg.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Register-blocked MatMul for the no-subgroup-matrix
@@ -13,31 +10,35 @@ enable f16;
 // accumulates through dot() and one step reads TM + TN words rather than
 // 4 * (TM + TN) scalars. A stores K contiguously and B stores N contiguously,
 // so each staging lane walks the axis its operand already has.
-{% set aR = source.aRank %}
-{% set bR = source.bRank %}
 {% set aBatchLen = aR - 2 %}
 {% set bBatchLen = bR - 2 %}
 {% set batchRank = aBatchLen %}
 {% set aTailStride = namespace(v=1) %}
-{% for j in range(1, aR) %}{% set aTailStride.v = aTailStride.v * source.aShape[j] %}{% endfor %}
 {% set bTailStride = namespace(v=1) %}
-{% for j in range(1, bR) %}{% set bTailStride.v = bTailStride.v * source.bShape[j] %}{% endfor %}
-{% if source.transBatchA %}{% set M = source.aShape[0] %}{% set K = source.aShape[aR-1] %}
-{% else %}{% set M = source.aShape[aR-2] %}{% set K = source.aShape[aR-1] %}{% endif %}
-{% set N = source.bShape[bR-1] %}
-{% if source.transBatchA %}{% set aMStride = aTailStride.v %}{% set aKStride = 1 %}
-{% else %}{% set aMStride = source.aShape[aR-1] %}{% set aKStride = 1 %}{% endif %}
-{% set bKStride = source.bShape[bR-1] %}{% set bNStride = 1 %}
-// Integer MatMul accumulates in its integer type because f32's 24-bit mantissa
-// rounds integer values above 2^24.
 {% set is_int = (scalar == "i32" or scalar == "u32") %}
 {% set accT = scalar if is_int else "f32" %}
-{% set outScalar = source.outScalar if source.outScalar is defined else scalar %}
-// f16 operands stay packed in workgroup memory and widen on shared load.
-// Accumulation remains f32 with the same FMA order, while shared traffic halves.
 {% set tileT = scalar if scalar == "f16" else accT %}
-{% set kTile = source.kTile if source.kTile is defined else 16 %}
 const M: u32 = {{ M }}u;
 const K: u32 = {{ K }}u;
 const N: u32 = {{ N }}u;
@@ -45,7 +46,7 @@ const A_M_STRIDE: u32 = {{ aMStride }}u;
 const A_K_STRIDE: u32 = {{ aKStride }}u;
 const B_K_STRIDE: u32 = {{ bKStride }}u;
 const B_N_STRIDE: u32 = {{ bNStride }}u;
-{% if is_int %}const ALPHA: {{ accT }} = {{ accT }}(1);{% else %}const ALPHA: f32 = f32({{ source.alpha }});{% endif %}
 // A 4x4 micro-tile over a 64x64 output tile reuses each staged operand across
 // four accumulators. It increases arithmetic work per load without the large
 // per-thread accumulator footprint of an 8x8 micro-tile.
@@ -54,6 +55,11 @@ 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
 const K_VECS: u32 = BK / 4u;
 var<workgroup> tileA: array<array<vec4<{{ tileT }}>, K_VECS>, BM>;   // A[m][k/4]
 var<workgroup> tileB: array<array<vec4<{{ tileT }}>, K_VECS>, BN>;   // B[n][k/4]
@@ -62,10 +68,10 @@ var<workgroup> tileB: array<array<vec4<{{ tileT }}>, K_VECS>, BN>;   // B[n][k/4
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
-{% set aStored = (aAxis + 1) if (source.transBatchA and aAxis >= 0) else aAxis %}
 {% set bStored = bAxis %}
-{% set aDim = source.aShape[aStored] %}
-{% set bDim = source.bShape[bStored] if bStored >= 0 else 1 %}
 {% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
 {% endfor %}
@@ -79,8 +85,12 @@ fn main(
   let li = lid.y * 16u + lid.x;
   let zOut = wg.z;
 {% if hasBatchCoord.value %}
   var zTmp = wg.z;
 {% endif %}
   var aBatchOff: u32 = 0u;
   var bBatchOff: u32 = 0u;
@@ -88,15 +98,15 @@ fn main(
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
-{% set aStored = (aAxis + 1) if (source.transBatchA and aAxis >= 0) else aAxis %}
 {% set bStored = bAxis %}
-{% set aDim = source.aShape[aStored] %}
-{% set bDim = source.bShape[bStored] if bStored >= 0 else 1 %}
 {% set outDim = aDim if aDim >= bDim else bDim %}
 {% set aStride = namespace(v=1) %}
-{% if aStored >= 0 and aDim != 1 %}{% for j in range(aStored + 1, aR) %}{% set aStride.v = aStride.v * source.aShape[j] %}{% endfor %}{% else %}{% set aStride.v = 0 %}{% endif %}
 {% set bStride = namespace(v=1) %}
-{% if bStored >= 0 and bDim != 1 %}{% for j in range(bStored + 1, bR) %}{% set bStride.v = bStride.v * source.bShape[j] %}{% endfor %}{% else %}{% set bStride.v = 0 %}{% endif %}
 {% if outDim > 1 %}
   let c{{ axis }} = zTmp % {{ outDim }}u;
   zTmp = zTmp / {{ outDim }}u;
@@ -110,9 +120,18 @@ fn main(
   var acc: array<{{ accT }}, TM * TN>;   // [ti*TN + tj] for the TMxTN micro-tile
   for (var i: u32 = 0u; i < TM * TN; i = i + 1u) { acc[i] = {{ accT }}(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. A's lanes walk K, which
     // it stores contiguously; B's walk N, which it stores contiguously.
     for (var idx: u32 = li; idx < BM * K_VECS; idx = idx + 256u) {
@@ -144,19 +163,24 @@ fn main(
       tileB[bc][br4] = bWord;
     }
     workgroupBarrier();
-{% set regIndent = "    " %}{% set regT = accT %}{{ 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] = vec4<{{ regT }}>(tileA[aRow + i][kv]); }
-{{ regIndent }}  for (var j: u32 = 0u; j < TN; j = j + 1u) { bv[j] = vec4<{{ regT }}>(tileB[bCol + j][kv]); }
-{{ 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();
   }

 {{ env.wgsl.resourceDeclarations }}
 // Register-blocked MatMul for the no-subgroup-matrix
 // accumulates through dot() and one step reads TM + TN words rather than
 // 4 * (TM + TN) scalars. A stores K contiguously and B stores N contiguously,
 // so each staging lane walks the axis its operand already has.
+{% set aR = aRank %}
+{% set bR = bRank %}
 {% set aBatchLen = aR - 2 %}
 {% set bBatchLen = bR - 2 %}
 {% set batchRank = aBatchLen %}
 {% set aTailStride = namespace(v=1) %}
+{% for j in range(1, aR) %}{% set aTailStride.v = aTailStride.v * aShape[j] %}{% endfor %}
 {% set bTailStride = namespace(v=1) %}
+{% for j in range(1, bR) %}{% set bTailStride.v = bTailStride.v * bShape[j] %}{% endfor %}
+{% if transBatchA %}{% set M = aShape[0] %}{% set K = aShape[aR-1] %}
+{% else %}{% set M = aShape[aR-2] %}{% set K = aShape[aR-1] %}{% endif %}
+{% set N = bShape[bR-1] %}
+{% if transBatchA %}{% set aMStride = aTailStride.v %}{% set aKStride = 1 %}
+{% else %}{% set aMStride = aShape[aR-1] %}{% set aKStride = 1 %}{% endif %}
+{% set bKStride = bShape[bR-1] %}{% set bNStride = 1 %}{% if bStorageStrides is defined %}{% set bKStride = bStorageStrides[bR-2] %}{% set bNStride = bStorageStrides[bR-1] %}{% endif %}
 {% set is_int = (scalar == "i32" or scalar == "u32") %}
+{% if is_int %}
+// Integer operands accumulate in their integer type, avoiding f32 rounding of
+// values outside the exact 24-bit significand range.
+{% endif %}
 {% set accT = scalar if is_int else "f32" %}
+{% set outScalar = outScalar if outScalar is defined else scalar %}
+{% set splitK = splitK if splitK is defined else 1 %}
+{% if scalar == "f16" %}
+// f16 operands stay packed in workgroup memory and widen on shared load;
+// accumulation remains f32.
+{% endif %}
 {% set tileT = scalar if scalar == "f16" else accT %}
+{% set kTile = kTile if kTile is defined else 16 %}
 const M: u32 = {{ M }}u;
 const K: u32 = {{ K }}u;
 const N: u32 = {{ N }}u;
 const A_K_STRIDE: u32 = {{ aKStride }}u;
 const B_K_STRIDE: u32 = {{ bKStride }}u;
 const B_N_STRIDE: u32 = {{ bNStride }}u;
+{% if is_int %}const ALPHA: {{ accT }} = {{ accT }}(1);{% else %}const ALPHA: f32 = f32({{ alpha }});{% endif %}
 // A 4x4 micro-tile over a 64x64 output tile reuses each staged operand across
 // four accumulators. It increases arithmetic work per load without the large
 // per-thread accumulator footprint of an 8x8 micro-tile.
 const BN: u32 = 64u;
 const TM: u32 = 4u;   // per-thread micro-tile rows
 const TN: u32 = 4u;   // per-thread micro-tile cols
+{% if splitK > 1 %}
+const SPLIT_K: u32 = {{ splitK }}u;
+const K_PER_SPLIT: u32 = {{ kPerSplit }}u;
+{% endif %}
 const K_VECS: u32 = BK / 4u;
 var<workgroup> tileA: array<array<vec4<{{ tileT }}>, K_VECS>, BM>;   // A[m][k/4]
 var<workgroup> tileB: array<array<vec4<{{ tileT }}>, K_VECS>, BN>;   // B[n][k/4]
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aStored = (aAxis + 1) if (transBatchA and aAxis >= 0) else aAxis %}
 {% set bStored = bAxis %}
+{% set aDim = aShape[aStored] %}
+{% set bDim = bShape[bStored] if bStored >= 0 else 1 %}
 {% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
 {% endfor %}
   let li = lid.y * 16u + lid.x;
   let zOut = wg.z;
+{% if splitK > 1 %}
+  let splitId = wg.z % SPLIT_K;
+{% else %}
 {% if hasBatchCoord.value %}
   var zTmp = wg.z;
+{% endif %}
 {% endif %}
   var aBatchOff: u32 = 0u;
   var bBatchOff: u32 = 0u;
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aStored = (aAxis + 1) if (transBatchA and aAxis >= 0) else aAxis %}
 {% set bStored = bAxis %}
+{% set aDim = aShape[aStored] %}
+{% set bDim = bShape[bStored] if bStored >= 0 else 1 %}
 {% set outDim = aDim if aDim >= bDim else bDim %}
 {% set aStride = namespace(v=1) %}
+{% if aStored >= 0 and aDim != 1 %}{% for j in range(aStored + 1, aR) %}{% set aStride.v = aStride.v * aShape[j] %}{% endfor %}{% else %}{% set aStride.v = 0 %}{% endif %}
 {% set bStride = namespace(v=1) %}
+{% if bStored >= 0 and bDim != 1 %}{% for j in range(bStored + 1, bR) %}{% set bStride.v = bStride.v * bShape[j] %}{% endfor %}{% else %}{% set bStride.v = 0 %}{% endif %}{% if bStorageStrides is defined and bStored >= 0 and bDim != 1 %}{% set bStride.v = bStorageStrides[bStored] %}{% endif %}
 {% if outDim > 1 %}
   let c{{ axis }} = zTmp % {{ outDim }}u;
   zTmp = zTmp / {{ outDim }}u;
   var acc: array<{{ accT }}, TM * TN>;   // [ti*TN + tj] for the TMxTN micro-tile
   for (var i: u32 = 0u; i < TM * TN; i = i + 1u) { acc[i] = {{ accT }}(0); }
+{% if splitK > 1 %}
+  let kStart = splitId * K_PER_SPLIT;
+  let numTiles = K_PER_SPLIT / BK;
+{% else %}
   let numTiles = (K + BK - 1u) / BK;
+{% endif %}
   for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+{% if splitK > 1 %}
+    let kBase = kStart + kt * BK;
+{% else %}
     let kBase = kt * BK;
+{% endif %}
     // Cooperative load: one vector word per lane per pass. A's lanes walk K, which
     // it stores contiguously; B's walk N, which it stores contiguously.
     for (var idx: u32 = li; idx < BM * K_VECS; idx = idx + 256u) {
       tileB[bc][br4] = bWord;
     }
     workgroupBarrier();
+{% set regT = accT %}{% filter indent(4, 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] = vec4<{{ regT }}>(tileA[aRow + i][kv]); }
+  for (var j: u32 = 0u; j < TN; j = j + 1u) { bv[j] = vec4<{{ regT }}>(tileB[bCol + j][kv]); }
+  for (var i: u32 = 0u; i < TM; i = i + 1u) {
+    for (var j: u32 = 0u; j < TN; j = j + 1u) {
+{% if regSequentialK is defined and regSequentialK %}{% for component in range(4) %}
+      acc[i * TN + j] = acc[i * TN + j] + av[i][{{ component }}] * bv[j][{{ component }}];
+{% endfor %}{% else %}      acc[i * TN + j] = acc[i * TN + j] + dot(av[i], bv[j]);
+{% endif %}
+    }
+  }
+}
+{% endfilter %}
     workgroupBarrier();
   }

build/webgpu/matmul-tiled-general.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Shared tiled matrix multiplication: Y = alpha * op(A) @ op(B), where op
@@ -10,8 +7,8 @@ enable f16;
 // K/M/N. transA/transB only change which stored stride the logical (m,k)/(k,n)
 // walk, using compiled stride constants here; the batch broadcast strides are
 // compiled the same way (a 0 literal means "broadcast / absent on that operand").
-{% set aR = source.aRank %}
-{% set bR = source.bRank %}
 {% set aVec = (aR == 1) %}
 {% set bVec = (bR == 1) %}
 {% set aBatchLen = (aR - 2) if aR >= 2 else 0 %}
@@ -22,28 +19,28 @@ enable f16;
  * axis is unchanged, and the remaining axes form the batch. This stride
  * permutation composes with the ordinary last-two-axis transpose. */
 {% set aTailStride = namespace(v=1) %}
-{% for j in range(1, aR) %}{% set aTailStride.v = aTailStride.v * source.aShape[j] %}{% endfor %}
 {% set bTailStride = namespace(v=1) %}
-{% for j in range(1, bR) %}{% set bTailStride.v = bTailStride.v * source.bShape[j] %}{% endfor %}
-{% if aVec %}{% set M = 1 %}{% set K = source.aShape[0] %}
-{% elif source.transBatchA and source.transA %}{% set M = source.aShape[aR-1] %}{% set K = source.aShape[0] %}
-{% elif source.transBatchA %}{% set M = source.aShape[0] %}{% set K = source.aShape[aR-1] %}
-{% elif source.transA %}{% set M = source.aShape[aR-1] %}{% set K = source.aShape[aR-2] %}
-{% else %}{% set M = source.aShape[aR-2] %}{% set K = source.aShape[aR-1] %}{% endif %}
 {% if bVec %}{% set N = 1 %}
-{% elif source.transBatchB and source.transB %}{% set N = source.bShape[0] %}
-{% elif source.transB %}{% set N = source.bShape[bR-2] %}
-{% else %}{% set N = source.bShape[bR-1] %}{% endif %}
 {% if aVec %}{% set aMStride = 0 %}{% set aKStride = 1 %}
-{% elif source.transBatchA and source.transA %}{% set aMStride = 1 %}{% set aKStride = aTailStride.v %}
-{% elif source.transBatchA %}{% set aMStride = aTailStride.v %}{% set aKStride = 1 %}
-{% elif source.transA %}{% set aMStride = 1 %}{% set aKStride = source.aShape[aR-1] %}
-{% else %}{% set aMStride = source.aShape[aR-1] %}{% set aKStride = 1 %}{% endif %}
 {% if bVec %}{% set bKStride = 1 %}{% set bNStride = 0 %}
-{% elif source.transBatchB and source.transB %}{% set bKStride = 1 %}{% set bNStride = bTailStride.v %}
-{% elif source.transBatchB %}{% set bKStride = bTailStride.v %}{% set bNStride = 1 %}
-{% elif source.transB %}{% set bKStride = 1 %}{% set bNStride = source.bShape[bR-1] %}
-{% else %}{% set bKStride = source.bShape[bR-1] %}{% set bNStride = 1 %}{% endif %}
 {% set is_int = (scalar == "i32" or scalar == "u32") %}
 {% set accT = scalar if is_int else "f32" %}
@@ -57,7 +54,7 @@ const B_K_STRIDE: u32 = {{ bKStride }}u;
 const B_N_STRIDE: u32 = {{ bNStride }}u;
 {% if is_int %}/* Integer matrix multiplication accumulates in the integer type. Widening
  * through f32 would round values above 2^24. Integer MatMul has alpha = 1. */
-const ALPHA: {{ accT }} = {{ accT }}(1);{% else %}const ALPHA: f32 = f32({{ source.alpha }});{% endif %}
 // 2x2 register-blocked tile: 16x16 threads each compute a 2x2 micro-tile, for a
 // 32x32 output tile per workgroup with K stepped in BK=16 chunks. Each loaded
 // shared-mem element feeds 2 FMAs, favoring register reuse in the inner loop.
@@ -72,10 +69,10 @@ var<workgroup> tileB: array<array<{{ tileT }}, 32>, 16>;
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
-{% set aStored = (aAxis + 1) if (source.transBatchA and aAxis >= 0) else aAxis %}
-{% set bStored = (bAxis + 1) if (source.transBatchB and bAxis >= 0) else bAxis %}
-{% set aDim = source.aShape[aStored] if aStored >= 0 else 1 %}
-{% set bDim = source.bShape[bStored] if bStored >= 0 else 1 %}
 {% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
 {% endfor %}
@@ -100,15 +97,15 @@ fn main(
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
-{% set aStored = (aAxis + 1) if (source.transBatchA and aAxis >= 0) else aAxis %}
-{% set bStored = (bAxis + 1) if (source.transBatchB and bAxis >= 0) else bAxis %}
-{% set aDim = source.aShape[aStored] if aStored >= 0 else 1 %}
-{% set bDim = source.bShape[bStored] if bStored >= 0 else 1 %}
 {% set outDim = aDim if aDim >= bDim else bDim %}
 {% set aStride = namespace(v=1) %}
-{% if aStored >= 0 and aDim != 1 %}{% for j in range(aStored + 1, aR) %}{% set aStride.v = aStride.v * source.aShape[j] %}{% endfor %}{% else %}{% set aStride.v = 0 %}{% endif %}
 {% set bStride = namespace(v=1) %}
-{% if bStored >= 0 and bDim != 1 %}{% for j in range(bStored + 1, bR) %}{% set bStride.v = bStride.v * source.bShape[j] %}{% endfor %}{% else %}{% set bStride.v = 0 %}{% endif %}
 {% if outDim > 1 %}
 {% if aStride.v != 0 or bStride.v != 0 %}
   let c{{ axis }} = zTmp % {{ outDim }}u;

 {{ env.wgsl.resourceDeclarations }}
 // Shared tiled matrix multiplication: Y = alpha * op(A) @ op(B), where op
 // K/M/N. transA/transB only change which stored stride the logical (m,k)/(k,n)
 // walk, using compiled stride constants here; the batch broadcast strides are
 // compiled the same way (a 0 literal means "broadcast / absent on that operand").
+{% set aR = aRank %}
+{% set bR = bRank %}
 {% set aVec = (aR == 1) %}
 {% set bVec = (bR == 1) %}
 {% set aBatchLen = (aR - 2) if aR >= 2 else 0 %}
  * axis is unchanged, and the remaining axes form the batch. This stride
  * permutation composes with the ordinary last-two-axis transpose. */
 {% set aTailStride = namespace(v=1) %}
+{% for j in range(1, aR) %}{% set aTailStride.v = aTailStride.v * aShape[j] %}{% endfor %}
 {% set bTailStride = namespace(v=1) %}
+{% for j in range(1, bR) %}{% set bTailStride.v = bTailStride.v * bShape[j] %}{% endfor %}
+{% if aVec %}{% set M = 1 %}{% set K = aShape[0] %}
+{% elif transBatchA and transA %}{% set M = aShape[aR-1] %}{% set K = aShape[0] %}
+{% elif transBatchA %}{% set M = aShape[0] %}{% set K = aShape[aR-1] %}
+{% elif transA %}{% set M = aShape[aR-1] %}{% set K = aShape[aR-2] %}
+{% else %}{% set M = aShape[aR-2] %}{% set K = aShape[aR-1] %}{% endif %}
 {% if bVec %}{% set N = 1 %}
+{% elif transBatchB and transB %}{% set N = bShape[0] %}
+{% elif transB %}{% set N = bShape[bR-2] %}
+{% else %}{% set N = bShape[bR-1] %}{% endif %}
 {% if aVec %}{% set aMStride = 0 %}{% set aKStride = 1 %}
+{% elif transBatchA and transA %}{% set aMStride = 1 %}{% set aKStride = aTailStride.v %}
+{% elif transBatchA %}{% set aMStride = aTailStride.v %}{% set aKStride = 1 %}
+{% elif transA %}{% set aMStride = 1 %}{% set aKStride = aShape[aR-1] %}
+{% else %}{% set aMStride = aShape[aR-1] %}{% set aKStride = 1 %}{% endif %}
 {% if bVec %}{% set bKStride = 1 %}{% set bNStride = 0 %}
+{% elif transBatchB and transB %}{% set bKStride = 1 %}{% set bNStride = bTailStride.v %}
+{% elif transBatchB %}{% set bKStride = bTailStride.v %}{% set bNStride = 1 %}
+{% elif transB %}{% set bKStride = 1 %}{% set bNStride = bShape[bR-1] %}
+{% else %}{% set bKStride = bShape[bR-1] %}{% set bNStride = 1 %}{% endif %}
 {% set is_int = (scalar == "i32" or scalar == "u32") %}
 {% set accT = scalar if is_int else "f32" %}
 const B_N_STRIDE: u32 = {{ bNStride }}u;
 {% if is_int %}/* Integer matrix multiplication accumulates in the integer type. Widening
  * through f32 would round values above 2^24. Integer MatMul has alpha = 1. */
+const ALPHA: {{ accT }} = {{ accT }}(1);{% else %}const ALPHA: f32 = f32({{ alpha }});{% endif %}
 // 2x2 register-blocked tile: 16x16 threads each compute a 2x2 micro-tile, for a
 // 32x32 output tile per workgroup with K stepped in BK=16 chunks. Each loaded
 // shared-mem element feeds 2 FMAs, favoring register reuse in the inner loop.
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aStored = (aAxis + 1) if (transBatchA and aAxis >= 0) else aAxis %}
+{% set bStored = (bAxis + 1) if (transBatchB and bAxis >= 0) else bAxis %}
+{% set aDim = aShape[aStored] if aStored >= 0 else 1 %}
+{% set bDim = bShape[bStored] if bStored >= 0 else 1 %}
 {% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
 {% endfor %}
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aStored = (aAxis + 1) if (transBatchA and aAxis >= 0) else aAxis %}
+{% set bStored = (bAxis + 1) if (transBatchB and bAxis >= 0) else bAxis %}
+{% set aDim = aShape[aStored] if aStored >= 0 else 1 %}
+{% set bDim = bShape[bStored] if bStored >= 0 else 1 %}
 {% set outDim = aDim if aDim >= bDim else bDim %}
 {% set aStride = namespace(v=1) %}
+{% if aStored >= 0 and aDim != 1 %}{% for j in range(aStored + 1, aR) %}{% set aStride.v = aStride.v * aShape[j] %}{% endfor %}{% else %}{% set aStride.v = 0 %}{% endif %}
 {% set bStride = namespace(v=1) %}
+{% if bStored >= 0 and bDim != 1 %}{% for j in range(bStored + 1, bR) %}{% set bStride.v = bStride.v * bShape[j] %}{% endfor %}{% else %}{% set bStride.v = 0 %}{% endif %}
 {% if outDim > 1 %}
 {% if aStride.v != 0 or bStride.v != 0 %}
   let c{{ axis }} = zTmp % {{ outDim }}u;

build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja CHANGED Viewed

@@ -1,17 +1,11 @@
-// GEMV specialization for y[N] = a[K] @ B[K, N]: each workgroup owns 32
-// consecutive vec4 column groups (128 output columns). Shared by MatMul and
-// matrix-shaped Einsum decode projections.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set OUT = outputBuffer if outputBuffer is defined else "c" %}
 const LANES: u32 = 32u;
-// SLICES is chosen from the workgroup COUNT, not from a fixed default: this
-// kernel launches ceilDiv(N, 128) workgroups, so a narrow-N projection starves
-// the device unless each workgroup carries more of the K reduction.
 const SLICES: u32 = {{ gemvSlices }}u;
 var<workgroup> partials: array<vec4<f32>, LANES * SLICES>;
@@ -37,6 +31,11 @@ fn main(
     for (var s = 1u; s < SLICES; s = s + 1u) {
       total = total + partials[s * LANES + lane];
     }
     {{ OUT }}[cg] = vec4<{{ T }}>(total);
   }
 }

+// GEMV specialization for y[N] = a[K] @ B[K, N]. Each workgroup owns 32
+// consecutive vec4 column groups, or 128 output columns.
 {{ env.wgsl.resourceDeclarations }}
 {% set OUT = outputBuffer if outputBuffer is defined else "c" %}
 const LANES: u32 = 32u;
+// SLICES partitions the K reduction across the workgroup's second dimension.
+// Thread zero of each column group combines the slice partials in index order.
 const SLICES: u32 = {{ gemvSlices }}u;
 var<workgroup> partials: array<vec4<f32>, LANES * SLICES>;
     for (var s = 1u; s < SLICES; s = s + 1u) {
       total = total + partials[s * LANES + lane];
     }
+{% if alphaScale is defined and alphaScale != 1 %}
+    // This specialization bakes the output multiplier into the shader and
+    // applies it after combining the K slices.
+    total = total * f32({{ alphaScale }});
+{% endif %}
     {{ OUT }}[cg] = vec4<{{ T }}>(total);
   }
 }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,23 +1,42 @@
 {
   "name": "com.microsoft.FusedMatMul",
-  "id": "_com_microsoft_fusedmatmul_webgpu_a7be42e",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "FKzhB1qpaKYuQ28OUS2fiiJaDroyLDXWubS0iFa3Fyw=",
-      "fused-matmul-subgroup-matrix.wgsl.jinja": "tZdeSq2o0dor9WP7R5Tf5Yzxm8DXHapNXdGci8SH4wI=",
-      "manifest.json": "TpmC/4RDrMS30inZSCKjBEw3wo+UAldsvOfUrpKuZn0=",
-      "matmul-subgroup-matrix-ext.wgsl.jinja": "W5PywIfrssv2VQhuCG1lu3LjC3R7VpxGRnkHwMX050Q=",
-      "matmul-tiled-general-reg.wgsl.jinja": "8E6pn6tg6t7LuoU6ktkjL5AZOgrd0HHkFRe8VJC7jwk=",
-      "matmul-tiled-general.wgsl.jinja": "pqYgOQCvHqeTthCEShE8dbN1GJjpFz3LD1K3sBYARMc=",
-      "matmul-vector-matrix-vec4.wgsl.jinja": "9syYZdAgoDnCn5rurCHP9tHtlwm1IXwKvPoj48hH3+E=",
-      "reduce-axis0-splitk-combine.wgsl.jinja": "uvQNNkj1ezvpKVhYyLUAAfGVLOk/i7VnAiiNSYYz2uk=",
-      "test.json": "Nv5BBqSF2s59CpK7a6Zmu1cc32XOn+PNz8J/MiLO0uo="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.FusedMatMul" }
 }

 {
   "name": "com.microsoft.FusedMatMul",
+  "id": "_com_microsoft_fusedmatmul_webgpu_8b2cb77",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "ATSeCpumMpLFypVxXuSfXbcCPQhVBBmfOof4PcQd89o=",
+      "fused-matmul-subgroup-matrix.wgsl.jinja": "0G1fPWFTk4qK1SiuAcriccOujenbWQfIgE0F0iXPlFo=",
+      "manifest.json": "pl4JzUzE2Z8hDvUoxFEkFv1qWvP/cMYQvOPN0s3asPs=",
+      "matmul-band-vec4.wgsl.jinja": "CN34bTiT4RnjmEoH/zPKP3a0tRpfYqP4l/X/bKCvHUk=",
+      "matmul-subgroup-matrix-ext.wgsl.jinja": "9hswKFk/g2HxvNxKwowsc7cJtEUGvrL53ISLi46Y3Co=",
+      "matmul-tiled-general-reg.wgsl.jinja": "Td1g9ghExH+uVbEhwxTEo1jkBpmkgKjBlkWGMD6JZ5I=",
+      "matmul-tiled-general.wgsl.jinja": "7vW3jelY9YQHhc92bE/yQZFQk4yeuSRby23ELji4rFs=",
+      "matmul-vector-matrix-vec4.wgsl.jinja": "TxVXDiRp6BURoOhTbh2ILDAAL9D9cVDfUhk8e5ViSfU=",
+      "reduce-axis0-splitk-combine.wgsl.jinja": "Yz1hjK55R/kndUrw3ugPqgPaOBwKmYupqZoVdJTHO5Q=",
+      "test.json": "CuQ/9/rI0a9gPgB8JwczBXvdF7PIuQoXpiFFKVmkeC4="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "subgroup_matrix_transbatch_b_f16": ["matmul-subgroup-matrix-ext.wgsl.jinja"],
+      "subgroup_matrix_transbatch_b_f32": ["matmul-subgroup-matrix-ext.wgsl.jinja"],
+      "f32_m1_gemv_vec4": ["matmul-vector-matrix-vec4.wgsl.jinja"],
+      "rank2_band_vec4_splitk": ["matmul-band-vec4.wgsl.jinja", "reduce-axis0-splitk-combine.wgsl.jinja"],
+      "rank2_band_vec4": ["matmul-band-vec4.wgsl.jinja"],
+      "rank2_band_vec4_f32_preferred": ["matmul-band-vec4.wgsl.jinja"],
+      "subgroup_matrix_splitk": ["matmul-subgroup-matrix-ext.wgsl.jinja", "reduce-axis0-splitk-combine.wgsl.jinja"],
+      "subgroup_matrix_tail_broadcast": ["matmul-subgroup-matrix-ext.wgsl.jinja"],
+      "subgroup_matrix": ["fused-matmul-subgroup-matrix.wgsl.jinja"],
+      "broadcast_rank4_tiled_reg": ["matmul-tiled-general-reg.wgsl.jinja"],
+      "plain_rank2_tiled_reg": ["matmul-tiled-general-reg.wgsl.jinja"],
+      "transbatch_a_tiled_reg": ["matmul-tiled-general-reg.wgsl.jinja"],
+      "tiled": ["matmul-tiled-general.wgsl.jinja"],
+      "transbatch_b_tiled_reg": ["matmul-tiled-general-reg.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja CHANGED Viewed

@@ -2,33 +2,121 @@
 // folds the segment partials and applies the selected reduction's final step.
 // Segments are folded in ascending order for deterministic results. This order
 // differs from the single-pass reduction but remains within the f32 tolerance.
-{% if source.outputF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-/* Exact max/min reduction identities. WGSL rejects infinity during constant
- * evaluation, so f32 identities are constructed at runtime from their
- * IEEE-754 bit patterns. */
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
   let stride = nwg.x * WG;
-  let start = (gid.y * nwg.x * WG) + gid.x;
   for (var col = start; col < params.cols; col = col + stride) {
     var total = 0.0;
     for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
       let p = partials[seg * params.cols + col];
       total = total + p;
     }
-{% if source.outputF16 %}
     y[col] = f16(total);
 {% else %}
     y[col] = total;
 {% endif %}
   }
 }

 // folds the segment partials and applies the selected reduction's final step.
 // Segments are folded in ascending order for deterministic results. This order
 // differs from the single-pass reduction but remains within the f32 tolerance.
+{% set addBias = addBias is defined and addBias %}
+{% set biasCols = biasCols | default(0) %}
+{% set intMode = intMode is defined and intMode %}
+{% set yv = "f16(" if outputF16 else "" %}
+{% set vy = ")" if outputF16 else "" %}
+{% if outputF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
+{% if addBias %}
+const BIAS_COLS: u32 = {{ biasCols }}u;
+{% endif %}
+{% if op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+const F32_MAX: f32 = 3.4028234663852886e38;
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{% elif op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("reduction_identity", op) }}
+{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
   let stride = nwg.x * WG;
+  let start = (gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG) + gid.x;
   for (var col = start; col < params.cols; col = col + stride) {
+{% if op == "logsumexp" %}
+    // Merge SPLIT (segMax, segSumExp) pairs stably; carry NaN / +Inf markers.
+    var nan_value = 0.0;
+    var has_nan = false;
+    var global_max = F32_MIN;
+    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
+      let nv = partials[(2u * SPLIT + seg) * params.cols + col];
+      if (nv != 0.0 || is_nan_f32(nv)) {
+        has_nan = true;
+        nan_value = nv;
+      }
+      global_max = max(global_max, partials[seg * params.cols + col]);
+    }
+    var sum = 0.0;
+    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
+      let seg_max = partials[seg * params.cols + col];
+      let seg_sum = partials[(SPLIT + seg) * params.cols + col];
+      sum = sum + seg_sum * exp(seg_max - global_max);
+    }
+    let has_positive_inf = global_max > F32_MAX;
+    let finite_or_inf = select(global_max + log(sum), global_max, has_positive_inf);
+    y[col] = {{ yv }}select(finite_or_inf, nan_value, has_nan){{ vy }};
+{% else %}
+{% if intMode %}
+{% if op == "prod" %}
+    var total = 1i;
+{% else %}
+    var total = 0i;
+{% endif %}
+{% else %}
+{% if op == "max" %}
+    var total = reduction_identity();
+{% elif op == "min" %}
+    var total = reduction_identity();
+{% elif op == "prod" %}
+    var total = 1.0;
+{% else %}
     var total = 0.0;
+{% endif %}
+{% endif %}
     for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
       let p = partials[seg * params.cols + col];
+{% if op == "max" or op == "min" %}
+      total = {{ op }}(total, p);
+{% elif op == "prod" %}
+      total = total * p;
+{% else %}
       total = total + p;
+{% endif %}
     }
+{% if addBias %}
+    total = total + f32(bias[col % BIAS_COLS]);
+{% endif %}
+{% if op == "l2" %}
+    y[col] = {{ yv }}sqrt(total){{ vy }};
+{% elif op == "logsum" %}
+    y[col] = {{ yv }}log(total){{ vy }};
+{% elif op == "mean" %}
+    y[col] = {{ yv }}total / f32(params.rows){{ vy }};
+{% else %}
+{% if outputF16 %}
     y[col] = f16(total);
 {% else %}
     y[col] = total;
+{% endif %}
+{% endif %}
 {% endif %}
   }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.FusedMatMul",
   "fixtureArrays": {
     "ort_float32_broadcast_rank3_by_rank4_output_Y": [1, 3, 5, 33, 43, 53, 5, 23, 41, 85, 111, 137, 9, 43, 77, 137, 179, 221],
     "ort_float32_rank3_by_rank2_output_Y": [20, 23, 26, 29, 56, 68, 80, 92, 92, 113, 134, 155, 128, 158, 188, 218],
@@ -297,12 +296,13 @@
       "name": "ort_float32_alpha_zero_outputs_zero",
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
-        "test": "FusedMatMulOpTest.DoubleTypeAlphaZero"
       },
       "attrs": { "alpha": 0 },
       "inputs": {
-        "A": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0] } },
-        "B": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [5.0, 6.0, 7.0, 8.0] } }
       },
       "outputs": {
         "Y": {
@@ -751,7 +751,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.02 } }
     },
     {
       "name": "aligned_f16_transB_alpha_64x32",
@@ -784,7 +784,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.011, "scale": 0.2 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [3, 7], "tolerance": 0.02 } }
     },
     {
       "name": "aligned_f16_transA_64x32",
@@ -801,7 +801,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03 } }
     },
     {
       "name": "aligned_f16_batched_plain_2x64x32x64",
@@ -817,7 +817,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [2, 64, 64], "tolerance": 0.03 } }
     },
     {
       "name": "f16_rank3_by_broadcast_rank3",
@@ -833,7 +833,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [2, 2, 4], "tolerance": 0.02 } }
     },
     {
       "name": "aligned_f16_transA_transB_alpha_64x32",
@@ -850,7 +850,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03 } }
     },
     {
       "name": "subgroup_matrix_m_tail_57_partial_block_f16",
@@ -866,7 +866,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.011, "scale": 0.2 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [57, 64], "tolerance": 0.03 } }
     },
     {
       "name": "subgroup_matrix_m_tail_33_alpha_scaled_f32",
@@ -931,7 +931,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [50, 64], "tolerance": 0.05 } }
     },
     {
       "name": "f32_decode_gemv_m1_k65_n68_vec4_compact",
@@ -989,7 +989,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [33, 66], "tolerance": 0.05 } }
     },
     {
       "name": "subgroup_matrix_broadcast_rank4x3_f16_compact",
@@ -1009,7 +1009,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 33, 64], "tolerance": 0.05 } }
     },
     {
       "name": "broadcast_rank4_tiled_reg_f16_compact",
@@ -1029,7 +1029,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 65, 67], "tolerance": 0.05 } }
     },
     {
       "name": "broadcast_rank4_tiled_reg_shared_f32_compact",
@@ -1069,7 +1069,7 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
         }
       },
-      "outputs": { "Y": { "dtype": "float16", "shape": [2, 65, 64], "tolerance": 0.05 } }
     },
     {
       "name": "transbatch_a_dense_m_tail_f32_compact",
@@ -1111,7 +1111,7 @@
     {
       "name": "subgroup_matrix_kn_tail_f16_offset_alpha_scale_lock",
       "provenance": {
-        "notes": "subgroup_matrix_kn_tail_f16_compact draws both operands from a zero-mean sinusoid, so the K=34 dot random-walks to 0.013 and the 0.05 absolute tolerance is 3.7x the whole output: alpha could be ignored, a K tile dropped, or Y halved and it still passes. Offsetting both operands makes every element approach alpha * K * aOffset * bOffset (~3.4), a value proportional to the contraction length, so the K=34 partial tile, the N=66 column tail and the alpha=0.5 epilogue are all under test on the tail-broadcast subgroup-matrix route."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
@@ -1131,7 +1131,7 @@
     {
       "name": "subgroup_matrix_broadcast_rank4x3_f16_offset_scale_lock",
       "provenance": {
-        "notes": "Well-conditioned sibling of subgroup_matrix_broadcast_rank4x3_f16_compact, whose zero-mean operands cancel to 0.021 under a 0.05 absolute tolerance (2.4x blind). Offsetting both operands drives each output to ~K * aOffset * bOffset, so the rank-4 by rank-3 batch broadcast now pins the per-batch B slice and the K=32 count instead of only the addressing pattern."
       },
       "attrs": { "alpha": 1 },
       "inputs": {
@@ -1151,7 +1151,7 @@
     {
       "name": "subgroup_matrix_a_batch_broadcast_rank4x3_f16",
       "provenance": {
-        "notes": "A-side batch broadcast: A carries dim 1 on the batch axis while B carries 2, so one A slice feeds both output batches. Every other rank-4 by rank-3 fixture broadcasts B (or matches both extents), leaving the A-broadcast stride - a zero A batch stride against a non-zero B one - unrendered. Operands are offset so each output lands near alpha * K * aOffset * bOffset and a swapped or dropped batch stride moves the result rather than cancelling."
       },
       "attrs": { "alpha": 1 },
       "inputs": {
@@ -1171,7 +1171,7 @@
     {
       "name": "broadcast_rank4_tiled_reg_f16_offset_alpha_scale_lock",
       "provenance": {
-        "notes": "broadcast_rank4_tiled_reg_f16_compact cancels to 0.015 against a 0.05 absolute tolerance, so the register-blocked rank-4 route was 3.4x blind to any uniform scale error - including a silently dropped alpha. Offsetting both operands makes each output ~alpha * K * aOffset * bOffset with K=33 (one full 32-wide k tile plus a 1-element tail), so a mis-counted k tile, a doubled accumulator or a missing alpha all move it far outside tolerance."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
@@ -1191,7 +1191,7 @@
     {
       "name": "transbatch_a_dense_m_tail_f16_offset_alpha_scale_lock",
       "provenance": {
-        "notes": "transbatch_a_dense_m_tail_f16_compact was the most blind fixture of this op: a 0.05 absolute tolerance over a 0.010 output, so a 4.8x uniform scale error passes. Offsetting both operands makes every element ~alpha * K * aOffset * bOffset, so the [M,batch,K] transBatchA stride, the K=32 contraction and the alpha=0.5 scale are pinned rather than just the row/column addressing."
       },
       "attrs": { "alpha": 0.5, "transBatchA": 1 },
       "inputs": {
@@ -1211,7 +1211,7 @@
     {
       "name": "aligned_f16_transA_transB_alpha_offset_scale_lock",
       "provenance": {
-        "notes": "aligned_f16_transA_transB_alpha_64x32 cancels to 0.11 under a 0.03 absolute tolerance, so a 27% scale error - alpha 0.5 applied twice, say - still passes. Offsetting both operands makes each output ~alpha * K * aOffset * bOffset, which keeps the doubly-transposed addressing under test while making the alpha epilogue and the K=32 count observable on both the subgroup-matrix and portable tiled routes."
       },
       "attrs": { "transA": 1, "transB": 1, "alpha": 0.5 },
       "inputs": {
@@ -1231,7 +1231,7 @@
     {
       "name": "f16_unaligned_3x5x7_offset_scale_lock",
       "provenance": {
-        "notes": "f16_unaligned_3x5x7 cancels to 0.15 under a 0.02 absolute tolerance (13% blind). The same 3x5 by 5x7 shape with offset operands lands at ~K * aOffset * bOffset, so the odd-K/odd-N scalar and tiled fallbacks - the routes small unaligned shapes actually take - now catch a dropped K element or a doubled tail."
       },
       "inputs": {
         "A": {
@@ -1250,7 +1250,7 @@
     {
       "name": "aligned_f16_plain_64x32x64_offset_scale_lock",
       "provenance": {
-        "notes": "aligned_f16_plain_64x32x64 cancels to 0.124 under a 0.02 absolute tolerance, so a 16% uniform scale error passes on the fully aligned f16 subgroup-matrix path (M=64, K=32, N=64 - no tail at all). Offsetting both operands makes each element ~K * aOffset * bOffset, so the aligned 32x64 tile now pins the k-block count and the scratch drain, not just the addressing."
       },
       "inputs": {
         "A": {
@@ -1269,7 +1269,7 @@
     {
       "name": "aligned_f16_batched_plain_2x64x32x64_offset_scale_lock",
       "provenance": {
-        "notes": "Batched twin of aligned_f16_plain_64x32x64_offset_scale_lock. aligned_f16_batched_plain_2x64x32x64 cancels to 0.124 under a 0.03 absolute tolerance (24% blind); with offset operands each element is ~K * aOffset * bOffset and the two batches carry different data, so the batch stride is pinned alongside the magnitude."
       },
       "inputs": {
         "A": {
@@ -1288,7 +1288,7 @@
     {
       "name": "subgroup_matrix_m_tail_57_partial_block_f16_offset_scale_lock",
       "provenance": {
-        "notes": "subgroup_matrix_m_tail_57_partial_block_f16 cancels to 0.062 under a 0.03 absolute tolerance, so the partial 8-row MMA block at M=57 (one full 32-row tile plus a 25-row remainder) was 49% blind to a scale error. With offset operands each element is ~K * aOffset * bOffset, so the guarded rows have to carry the same magnitude as the full ones - a tail block that accumulates one k step short, or drains a stale scratch slot, no longer looks like noise."
       },
       "inputs": {
         "A": {
@@ -1307,7 +1307,7 @@
     {
       "name": "aligned_f16_transA_64x32_offset_scale_lock",
       "provenance": {
-        "notes": "aligned_f16_transA_64x32 cancels to 0.053 under a 0.03 absolute tolerance - 57% blind, the loosest of the transpose fixtures. transA alone (transB = 0) is its own render of the operand loader, and with offset operands each element is ~K * aOffset * bOffset, so the transposed A stride is pinned against a k-count error rather than only against a wrong-element read."
       },
       "attrs": { "transA": 1 },
       "inputs": {
@@ -1327,7 +1327,7 @@
     {
       "name": "transA_transB_subgroup_matrix_m_tail_50_f16_offset_scale_lock",
       "provenance": {
-        "notes": "Combines both transposes with an M tail (50 = 32 + 18). Offset operands make each element approximately K * aOffset * bOffset, checking the guarded rows of the doubly-transposed load for magnitude as well as placement; the zero-mean companion cancels to 0.099 under a 0.05 absolute tolerance and cannot provide that scale lock."
       },
       "attrs": { "transA": 1, "transB": 1, "alpha": 1 },
       "inputs": {
@@ -1467,6 +1467,26 @@
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [16, 128], "tolerance": 0.0002 } }
     },
     {
       "name": "subgroup_matrix_batched_transB_small_m_f16",
       "attrs": { "transB": 1, "alpha": 0.25 },
@@ -1500,6 +1520,240 @@
         }
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [8, 64], "tolerance": 0.005 } }
     }
   ]
 }

 {
   "fixtureArrays": {
     "ort_float32_broadcast_rank3_by_rank4_output_Y": [1, 3, 5, 33, 43, 53, 5, 23, 41, 85, 111, 137, 9, 43, 77, 137, 179, 221],
     "ort_float32_rank3_by_rank2_output_Y": [20, 23, 26, 29, 56, 68, 80, 92, 92, 113, 134, 155, 128, 158, 188, 218],
       "name": "ort_float32_alpha_zero_outputs_zero",
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.DoubleTypeAlphaZero",
+        "notes": "Diverges from the upstream test's inputs (inputs.A values [1.0, 2.0, 3.0, 4.0] -> constant 2.0; inputs.B values [5.0, 6.0, 7.0, 8.0] -> constant 3.0); the expected output is recomputed by the CPU reference for the new inputs. A zero alpha scales the whole product away, so no operand value can reach the result and both operands are uniform fills."
       },
       "attrs": { "alpha": 0 },
       "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "constant", "value": 2.0 } },
+        "B": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "constant", "value": 3.0 } }
       },
       "outputs": {
         "Y": {
           "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.002 } }
     },
     {
       "name": "aligned_f16_transB_alpha_64x32",
           "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.011, "scale": 0.2 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [3, 7], "tolerance": 0.003 } }
     },
     {
       "name": "aligned_f16_transA_64x32",
           "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.001 } }
     },
     {
       "name": "aligned_f16_batched_plain_2x64x32x64",
           "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 64, 64], "tolerance": 0.002 } }
     },
     {
       "name": "f16_rank3_by_broadcast_rank3",
           "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 2, 4], "tolerance": 0.002 } }
     },
     {
       "name": "aligned_f16_transA_transB_alpha_64x32",
           "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.002 } }
     },
     {
       "name": "subgroup_matrix_m_tail_57_partial_block_f16",
           "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.011, "scale": 0.2 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [57, 64], "tolerance": 0.001 } }
     },
     {
       "name": "subgroup_matrix_m_tail_33_alpha_scaled_f32",
           "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [50, 64], "tolerance": 0.002 } }
     },
     {
       "name": "f32_decode_gemv_m1_k65_n68_vec4_compact",
           "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [33, 66], "tolerance": 0.0004 } }
     },
     {
       "name": "subgroup_matrix_broadcast_rank4x3_f16_compact",
           "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 33, 64], "tolerance": 0.0005 } }
     },
     {
       "name": "broadcast_rank4_tiled_reg_f16_compact",
           "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 65, 67], "tolerance": 0.0004 } }
     },
     {
       "name": "broadcast_rank4_tiled_reg_shared_f32_compact",
           "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
         }
       },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 65, 64], "tolerance": 0.0004 } }
     },
     {
       "name": "transbatch_a_dense_m_tail_f32_compact",
     {
       "name": "subgroup_matrix_kn_tail_f16_offset_alpha_scale_lock",
       "provenance": {
+        "notes": "Offset float16 operands keep each output near `alpha * K * aOffset * bOffset` (about 3.4), making the K=34 reduction tail, N=66 column tail, and `alpha = 0.5` epilogue observable on the subgroup-matrix route."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
     {
       "name": "subgroup_matrix_broadcast_rank4x3_f16_offset_scale_lock",
       "provenance": {
+        "notes": "Offset operands keep outputs near `K * aOffset * bOffset`, making the per-batch B slice and K=32 contraction observable in a rank-4 by rank-3 broadcast."
       },
       "attrs": { "alpha": 1 },
       "inputs": {
     {
       "name": "subgroup_matrix_a_batch_broadcast_rank4x3_f16",
       "provenance": {
+        "notes": "A has batch extent 1 while B has extent 2, so one A slice feeds both output batches. Offset operands keep the expected magnitude nonzero, exposing a swapped or nonzero A batch stride."
       },
       "attrs": { "alpha": 1 },
       "inputs": {
     {
       "name": "broadcast_rank4_tiled_reg_f16_offset_alpha_scale_lock",
       "provenance": {
+        "notes": "Offset operands keep outputs near `alpha * K * aOffset * bOffset` with K=33, making the one-element reduction tail and `alpha` multiplier observable on the register-blocked rank-4 route."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
     {
       "name": "transbatch_a_dense_m_tail_f16_offset_alpha_scale_lock",
       "provenance": {
+        "notes": "Offset operands make each output proportional to `alpha * K`, exposing the [M, batch, K] transBatchA stride, K=32 contraction, and `alpha = 0.5` scale."
       },
       "attrs": { "alpha": 0.5, "transBatchA": 1 },
       "inputs": {
     {
       "name": "aligned_f16_transA_transB_alpha_offset_scale_lock",
       "provenance": {
+        "notes": "Offset operands keep the doubly transposed output proportional to `alpha * K`, making the K=32 contraction and `alpha = 0.5` epilogue observable on subgroup-matrix and portable tiled routes."
       },
       "attrs": { "transA": 1, "transB": 1, "alpha": 0.5 },
       "inputs": {
     {
       "name": "f16_unaligned_3x5x7_offset_scale_lock",
       "provenance": {
+        "notes": "Offset operands in a 3x5 by 5x7 multiply keep outputs proportional to K, exposing dropped reduction elements or doubled tails on the unaligned scalar and tiled routes."
       },
       "inputs": {
         "A": {
     {
       "name": "aligned_f16_plain_64x32x64_offset_scale_lock",
       "provenance": {
+        "notes": "Offset operands keep each fully aligned M=64, K=32, N=64 output proportional to K, making the subgroup-matrix reduction count and scratch drain observable."
       },
       "inputs": {
         "A": {
     {
       "name": "aligned_f16_batched_plain_2x64x32x64_offset_scale_lock",
       "provenance": {
+        "notes": "Two batches carry distinct offset operands, keeping outputs proportional to K and making both the batch stride and aligned subgroup-matrix reduction count observable."
       },
       "inputs": {
         "A": {
     {
       "name": "subgroup_matrix_m_tail_57_partial_block_f16_offset_scale_lock",
       "provenance": {
+        "notes": "M=57 leaves 25 rows after one full 32-row tile. Offset operands require tail rows to match the full rows' expected magnitude, exposing a short reduction or stale scratch value."
       },
       "inputs": {
         "A": {
     {
       "name": "aligned_f16_transA_64x32_offset_scale_lock",
       "provenance": {
+        "notes": "With only A transposed, offset operands keep each output proportional to K and expose both a transposed-A stride error and an incorrect reduction count."
       },
       "attrs": { "transA": 1 },
       "inputs": {
     {
       "name": "transA_transB_subgroup_matrix_m_tail_50_f16_offset_scale_lock",
       "provenance": {
+        "notes": "Both operands are transposed and M=50 leaves an 18-row tail. Offset operands make the guarded tail rows' magnitude and placement independently observable."
       },
       "attrs": { "transA": 1, "transB": 1, "alpha": 1 },
       "inputs": {
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [16, 128], "tolerance": 0.0002 } }
     },
+    {
+      "name": "band_vec4_alpha_scaled_m8_k256_n512",
+      "provenance": {
+        "notes": "Route lock for the few-row band without a K split: eight rows over four 128-column groups, alpha folded into the vec4 store."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [8, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.13, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [256, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.19, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 512], "tolerance": 0.0001 } }
+    },
     {
       "name": "subgroup_matrix_batched_transB_small_m_f16",
       "attrs": { "transB": 1, "alpha": 0.25 },
         }
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [8, 64], "tolerance": 0.005 } }
+    },
+    {
+      "name": "f32_decode_gemv_m1_k65_n68_alpha_half",
+      "provenance": {
+        "notes": "A compact M=1 GEMV with alpha=0.5 exercises a non-unit multiplier folded into the store as a baked constant. The expected output verifies that the multiplier is applied exactly once."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [1, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [65, 68],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [1, 68], "tolerance": 0.0002 } }
+    },
+    {
+      "name": "f16_rank4_by_rank2_shared_weight_m33_k34_n66",
+      "provenance": {
+        "notes": "Shared rank-2 B across both rank-4 batch axes; covers subgroup-matrix admission, partial tiles and its small-M fallback without changing accumulation precision."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 33, 34],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [34, 66],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 33, 66], "tolerance": 0.001 } }
+    },
+    {
+      "name": "f16_rank4_by_rank2_shared_weight_m8_k32_n64",
+      "provenance": {
+        "notes": "Shared rank-2 B across both rank-4 batch axes; covers subgroup-matrix admission, partial tiles and its small-M fallback without changing accumulation precision."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 8, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 8, 64], "tolerance": 0.001 } }
+    },
+    {
+      "name": "f16_rank4_by_rank2_shared_weight_m1_k32_n64",
+      "provenance": {
+        "notes": "Shared rank-2 B across both rank-4 batch axes; covers subgroup-matrix admission, partial tiles and its small-M fallback without changing accumulation precision."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 1, 64], "tolerance": 0.001 } }
+    },
+    {
+      "name": "transbatch_b_aligned_occupancy_float16",
+      "attrs": { "transBatchB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [128, 2, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.07, "offset": 0.3 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 128, 512], "tolerance": 0.001, "relTolerance": 0.001 } }
+    },
+    {
+      "name": "transbatch_b_aligned_occupancy_float32",
+      "attrs": { "transBatchB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [128, 2, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.07, "offset": 0.3 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 128, 512], "tolerance": 0.0001, "relTolerance": 0.00002 } }
+    },
+    {
+      "name": "transbatch_b_portable_reg_aligned_float16",
+      "attrs": { "transBatchB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [4, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [128, 4, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.07, "offset": 0.3 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 128, 512], "tolerance": 0.001, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Portable register tile transBatchB: physical B[K,batch,N], natural tile count reaches the shared occupancy gate; negative alpha and positive inputs avoid cancellation."
+      }
+    },
+    {
+      "name": "transbatch_b_portable_reg_tails_float16",
+      "attrs": { "transBatchB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [3, 129, 131],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [131, 3, 513],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.07, "offset": 0.3 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [3, 129, 513], "tolerance": 0.001, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Portable register tile transBatchB: physical B[K,batch,N], natural tile count reaches the shared occupancy gate; negative alpha and positive inputs avoid cancellation."
+      }
+    },
+    {
+      "name": "transbatch_b_portable_reg_aligned_float32",
+      "attrs": { "transBatchB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [4, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [128, 4, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.07, "offset": 0.3 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 128, 512], "tolerance": 0.0001, "relTolerance": 0.00002 } },
+      "provenance": {
+        "notes": "Portable register tile transBatchB: physical B[K,batch,N], natural tile count reaches the shared occupancy gate; negative alpha and positive inputs avoid cancellation."
+      }
+    },
+    {
+      "name": "transbatch_b_portable_reg_tails_float32",
+      "attrs": { "transBatchB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 129, 131],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [131, 3, 513],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.07, "offset": 0.3 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [3, 129, 513], "tolerance": 0.0001, "relTolerance": 0.00002 } },
+      "provenance": {
+        "notes": "Portable register tile transBatchB: physical B[K,batch,N], natural tile count reaches the shared occupancy gate; negative alpha and positive inputs avoid cancellation."
+      }
+    },
+    {
+      "name": "f32_band_preferred_m4_k2048_n4096",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [4, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.1, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2048, 4096],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.041, "scale": 0.1, "offset": 0.03 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 4096], "tolerance": 0.0001, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Natural f32 preferred single-band geometry with nonuniform operands and alpha0.5. N4096 keeps the existing single-band path eligible at deep K; N2048 would select the split-band family."
+      }
+    },
+    {
+      "name": "f32_band_preferred_m16_k2560_n4096",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [16, 2560],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.1, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2560, 4096],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.041, "scale": 0.1, "offset": 0.03 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [16, 4096], "tolerance": 0.0001, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Natural f32 preferred single-band geometry with nonuniform operands and alpha0.5. N4096 keeps the existing single-band path eligible at deep K; N2048 would select the split-band family."
+      }
     }
   ]
 }