sync 6fdf6301e2bb

Browse files

Files changed (12) hide show

README.md +7 -2
build/webgpu/bench.json +305 -16
build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja +44 -16
build/webgpu/manifest.json +234 -143
build/webgpu/matmul-band-vec4.wgsl.jinja +2 -6
build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja +36 -19
build/webgpu/matmul-tiled-general-reg.wgsl.jinja +37 -21
build/webgpu/matmul-tiled-general.wgsl.jinja +17 -15
build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja +16 -2
build/webgpu/metadata.json +17 -14
build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja +4 -31
build/webgpu/test.json +694 -29

README.md CHANGED Viewed

@@ -51,12 +51,17 @@ Default values (overridable per request):
 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.
@@ -69,7 +74,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 - [`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)
@@ -81,7 +86,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 ## 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.

 One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `broadcast_transb_tiled_reg` — Register-blocked broadcast product with physically transposed B. Reuses the shared batch-addressing tile, keeps f32 accumulation, and preserves scalar K order for f16. Low tile count and excessive padding demote this otherwise correct path.
+- `broadcast_transb_subgroup_matrix_f16` — Broadcast transposed-B product using a supported 8x8x8 subgroup-matrix configuration with f32 accumulation. Logical shapes and physical B strides share the existing matrix engine; insufficient output tiles or excessive padding retain the generic tile.
+- `broadcast_transb_subgroup_matrix_f32` — Broadcast transposed-B product using a supported 8x8x8 subgroup-matrix configuration with f32 accumulation. Logical shapes and physical B strides share the existing matrix engine; insufficient output tiles or excessive padding retain the generic tile.
 - `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.
+- `m1_gemv_vec4` — Vector-by-matrix specialization for a single output row: each workgroup owns 32 consecutive vec4 column groups and partitions the reduction across the workgroup's second dimension. The accumulator stays float32 for both tensor types.
 - `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.
+- `subgroup_matrix` — Subgroup-matrix `Y = alpha * op(A) @ op(B)` over dense batches with float32 accumulation. An output width that is not a multiple of the 64-wide column tile switches the trailing tile to guarded addressing: its B columns clamp to N - 1 and its stores drop every column at or past N. Yields the shape when the padded column ratio exceeds its tunable ceiling.
 - `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.
 - [`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 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)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/bench.json CHANGED Viewed

@@ -56,7 +56,62 @@
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * 384 * 400 * 1536" }] }
     },
     {
-      "name": "fusedmatmul-f16-aligned-n512-512x2048x512-healthy",
       "preset": "smoke",
       "attrs": { "alpha": 1 },
       "vars": { "M": 512, "K": 2048, "N": 512 },
@@ -163,7 +218,7 @@
       "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 },
@@ -175,7 +230,7 @@
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * args.K" }] }
     },
     {
-      "name": "fusedmatmul-f16-transbatch-a-forces-tiled-8x512x2048x512-stress",
       "preset": "stress",
       "attrs": { "alpha": 1, "transBatchA": 1 },
       "vars": { "batch": 8, "M": 512, "K": 2048, "N": 512 },
@@ -190,7 +245,7 @@
       "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 },
@@ -202,7 +257,7 @@
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 512 * 512 * 2048" }] }
     },
     {
-      "name": "fusedmatmul-f16-broadcast-batch-forces-tiled-1x8x512x2048x512-stress",
       "preset": "stress",
       "attrs": { "alpha": 1 },
       "inputs": {
@@ -305,7 +360,7 @@
       "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 },
@@ -319,9 +374,7 @@
     {
       "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": {
@@ -334,9 +387,7 @@
     {
       "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": {
@@ -363,7 +414,7 @@
       },
       "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)" }] }
@@ -385,7 +436,7 @@
       },
       "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)" }] }
@@ -407,7 +458,7 @@
       },
       "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)" }] }
@@ -429,7 +480,7 @@
       },
       "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)" }] }
@@ -499,6 +550,244 @@
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [16, 4096], "dist": "empty" } },
       "bench": { "metrics": [{ "type": "gflops", "value": 1073741824 }] }
     }
   ]
 }

       "bench": { "metrics": [{ "type": "gflops", "value": "2 * 384 * 400 * 1536" }] }
     },
     {
+      "name": "fusedmatmul-f16-deep-skew-m128-k2048-n2048",
+      "preset": "model",
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": { "shape": [128, 2048], "dtype": "float16", "dist": "normal", "seed": 596, "scale": 0.1 },
+        "B": { "shape": [2048, 2048], "dtype": "float16", "dist": "normal", "seed": 597, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [128, 2048], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 128 * 2048 * 2048" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-deep-skew-m1024-k2048-n256",
+      "preset": "model",
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": { "shape": [1024, 2048], "dtype": "float16", "dist": "normal", "seed": 598, "scale": 0.1 },
+        "B": { "shape": [2048, 256], "dtype": "float16", "dist": "normal", "seed": 599, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [1024, 256], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 1024 * 256 * 2048" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-depth-k64-512x64x512",
+      "preset": "model",
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": { "shape": [512, 64], "dtype": "float16", "dist": "normal", "seed": 590, "scale": 0.1 },
+        "B": { "shape": [64, 512], "dtype": "float16", "dist": "normal", "seed": 591, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [512, 512], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 512 * 512 * 64" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-depth-k512-512x512x512",
+      "preset": "model",
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": { "shape": [512, 512], "dtype": "float16", "dist": "normal", "seed": 592, "scale": 0.1 },
+        "B": { "shape": [512, 512], "dtype": "float16", "dist": "normal", "seed": 593, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [512, 512], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 512 * 512 * 512" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-depth-k1024-512x1024x512",
+      "preset": "model",
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": { "shape": [512, 1024], "dtype": "float16", "dist": "normal", "seed": 594, "scale": 0.1 },
+        "B": { "shape": [1024, 512], "dtype": "float16", "dist": "normal", "seed": 595, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [512, 512], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 512 * 512 * 1024" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-aligned-n512-512x2048x512-control",
       "preset": "smoke",
       "attrs": { "alpha": 1 },
       "vars": { "M": 512, "K": 2048, "N": 512 },
       "name": "fusedmatmul-f16-rank4-by-rank2-shared-weight-b2h8-m512-k2048-n512-pathology",
       "preset": "stress",
       "provenance": {
+        "notes": "Measures FusedMatMul over a batched (2x8) projection sharing one rank-2 [2048,512] weight, with M=512, K=2048, N=512 (float16)."
       },
       "attrs": { "alpha": 1 },
       "vars": { "dtype": "float16", "M": 512, "K": 2048, "N": 512 },
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * args.K" }] }
     },
     {
+      "name": "fusedmatmul-f16-transbatch-a-tiled-8x512x2048x512-stress",
       "preset": "stress",
       "attrs": { "alpha": 1, "transBatchA": 1 },
       "vars": { "batch": 8, "M": 512, "K": 2048, "N": 512 },
       "name": "fusedmatmul-f16-transbatch-b-8x512x2048x512-pathology",
       "preset": "stress",
       "provenance": {
+        "notes": "Measures FusedMatMul over a batch=8 projection with transposed-batch B (physical [K=2048, batch=8, N=512]): M=512, K=2048, N=512 (float16)."
       },
       "attrs": { "alpha": 1, "transBatchB": 1 },
       "vars": { "batch": 8, "M": 512, "K": 2048, "N": 512 },
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 512 * 512 * 2048" }] }
     },
     {
+      "name": "fusedmatmul-f16-broadcast-batch-tiled-1x8x512x2048x512-stress",
       "preset": "stress",
       "attrs": { "alpha": 1 },
       "inputs": {
       "name": "fusedmatmul-f32-transbatch-b-8x512x2048x512-pathology",
       "preset": "stress",
       "provenance": {
+        "notes": "Measures FusedMatMul over a batch=8 projection with transposed-batch B (physical [K=2048, batch=8, N=512]): M=512, K=2048, N=512 (float32)."
       },
       "attrs": { "alpha": 1, "transBatchB": 1 },
       "vars": { "batch": 8, "M": 512, "K": 2048, "N": 512 },
     {
       "name": "fusedmatmul-float16-transbatch-b-k128-workgroups64",
       "preset": "stress",
+      "provenance": { "notes": "Aligned interleaved B layout with a long reduction and 64 output tile groups." },
       "attrs": { "alpha": 0.5, "transBatchB": 1 },
       "vars": { "batch": 2, "M": 128, "K": 128, "N": 512 },
       "inputs": {
     {
       "name": "fusedmatmul-float32-transbatch-b-k128-workgroups64",
       "preset": "stress",
+      "provenance": { "notes": "Aligned interleaved B layout with a long reduction and 64 output tile groups." },
       "attrs": { "alpha": 0.5, "transBatchB": 1 },
       "vars": { "batch": 2, "M": 128, "K": 128, "N": 512 },
       "inputs": {
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [4, 128, 512] } },
       "provenance": {
+        "notes": "Measures FusedMatMul over a fully tile-aligned M=128, K=128, N=512 shape with transposed-batch B (float16, batch=4)."
       },
       "preset": "smoke",
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, 2)" }] }
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [3, 129, 513] } },
       "provenance": {
+        "notes": "Measures FusedMatMul over a transposed-batch-B shape with partial tiles: M=129, K=131, N=513 (float16, batch=3)."
       },
       "preset": "smoke",
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, 2)" }] }
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [4, 128, 512] } },
       "provenance": {
+        "notes": "Measures FusedMatMul over a fully tile-aligned M=128, K=128, N=512 shape with transposed-batch B (float32, batch=4)."
       },
       "preset": "smoke",
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, 2)" }] }
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [3, 129, 513] } },
       "provenance": {
+        "notes": "Measures FusedMatMul over a transposed-batch-B shape with partial tiles: M=129, K=131, N=513 (float32, batch=3)."
       },
       "preset": "smoke",
       "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, 2)" }] }
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [16, 4096], "dist": "empty" } },
       "bench": { "metrics": [{ "type": "gflops", "value": 1073741824 }] }
+    },
+    {
+      "name": "fusedmatmul-transb-tails-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [2, 1, 129, 65], "dist": "normal", "seed": 7302, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [3, 129, 65], "dist": "normal", "seed": 7303, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 129, 129] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Measures FusedMatMul over a rank-4 by rank-3 broadcast (batch dims 2x1 against 3) with transposed B: M=129, K=65, N=129 (float16)."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-rank3x2-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [4, 256, 128], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [512, 128], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 256, 512] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Measures FusedMatMul over a rank-3 by rank-2 broadcast with transposed B: M=256, K=128, N=512, batch=4 (float16)."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-rank5x3-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [2, 1, 2, 128, 32], "dist": "normal", "seed": 7306, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [2, 128, 32], "dist": "normal", "seed": 7307, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 1, 2, 128, 128] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Measures FusedMatMul over a rank-5 by rank-3 broadcast with transposed B: M=128, K=32, N=128, batch dims 2x1x2 (float16)."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-low_tiles-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [1, 64, 32], "dist": "normal", "seed": 7308, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [64, 32], "dist": "normal", "seed": 7309, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 64, 64] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Measures FusedMatMul over a rank-3 by rank-2 broadcast with transposed B: M=64, K=32, N=64, batch=1 (float16)."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-large-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 8, 512, 64], "dist": "normal", "seed": 7300, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [8, 512, 64], "dist": "normal", "seed": 7301, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 8, 512, 512] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Measures FusedMatMul over a rank-4 by rank-3 broadcast with transposed B: M=512, K=64, N=512, batch dims 2x8 (float32)."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-tails-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 1, 129, 65], "dist": "normal", "seed": 7302, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [3, 129, 65], "dist": "normal", "seed": 7303, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 129, 129] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Measures FusedMatMul over a rank-4 by rank-3 broadcast (batch dims 2x1 against 3) with transposed B: M=129, K=65, N=129 (float32)."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-rank3x2-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [4, 256, 128], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [512, 128], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 256, 512] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Measures FusedMatMul over a rank-3 by rank-2 broadcast with transposed B: M=256, K=128, N=512, batch=4 (float32)."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-rank5x3-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 1, 2, 128, 32], "dist": "normal", "seed": 7306, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [2, 128, 32], "dist": "normal", "seed": 7307, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 1, 2, 128, 128] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Measures FusedMatMul over a rank-5 by rank-3 broadcast with transposed B: M=128, K=32, N=128, batch dims 2x1x2 (float32)."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-low_tiles-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [1, 64, 32], "dist": "normal", "seed": 7308, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [64, 32], "dist": "normal", "seed": 7309, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [1, 64, 64] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Measures FusedMatMul over a rank-3 by rank-2 broadcast with transposed B: M=64, K=32, N=64, batch=1 (float32)."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-matrix_tile_floor-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [2, 256, 64], "dist": "normal", "seed": 7320, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [256, 64], "dist": "normal", "seed": 7321, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 256, 256] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Profitability-boundary probe for output tile count and padded arithmetic, with a shared transposed B matrix."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-register_tile_floor-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [4, 256, 64], "dist": "normal", "seed": 7322, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [256, 64], "dist": "normal", "seed": 7323, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 256, 256] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Profitability-boundary probe for output tile count and padded arithmetic, with a shared transposed B matrix."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-padding_inside-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [8, 256, 128], "dist": "normal", "seed": 7324, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [129, 128], "dist": "normal", "seed": 7325, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 256, 129] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Profitability-boundary probe for output tile count and padded arithmetic, with a shared transposed B matrix."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-padding_cross_engine-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [8, 256, 65], "dist": "normal", "seed": 7326, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [129, 65], "dist": "normal", "seed": 7327, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 256, 129] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Profitability-boundary probe for output tile count and padded arithmetic, with a shared transposed B matrix."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-matrix_tile_floor-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 256, 64], "dist": "normal", "seed": 7320, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [256, 64], "dist": "normal", "seed": 7321, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 256, 256] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Profitability-boundary probe for output tile count and padded arithmetic, with a shared transposed B matrix."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-register_tile_floor-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [4, 256, 64], "dist": "normal", "seed": 7322, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [256, 64], "dist": "normal", "seed": 7323, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 256, 256] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Profitability-boundary probe for output tile count and padded arithmetic, with a shared transposed B matrix."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-padding_inside-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [8, 256, 128], "dist": "normal", "seed": 7324, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [129, 128], "dist": "normal", "seed": 7325, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 256, 129] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Profitability-boundary probe for output tile count and padded arithmetic, with a shared transposed B matrix."
+      }
+    },
+    {
+      "name": "fusedmatmul-transb-padding_cross_engine-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [8, 256, 65], "dist": "normal", "seed": 7326, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [129, 65], "dist": "normal", "seed": 7327, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 256, 129] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, -1)" }] },
+      "provenance": {
+        "notes": "Profitability-boundary probe for output tile count and padded arithmetic, with a shared transposed B matrix."
+      }
     }
   ]
 }

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

@@ -1,9 +1,17 @@
 // 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;
@@ -11,27 +19,22 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = fScalar %}
 {% set accScalar = "f32" %}
-const M: u32 = {{ M }}u;
 const K: u32 = {{ K }}u;
 const N: u32 = {{ N }}u;
 {% if transBatchA %}
 const BATCH_COUNT: u32 = {{ batchCount }}u;
 const A_BATCH_STRIDE: u32 = K;
 const A_M_STRIDE: u32 = BATCH_COUNT * K;
-{% else %}
-const A_BATCH_STRIDE: u32 = M * K;
-{% if not transA %}
 const A_M_STRIDE: u32 = K;
 {% endif %}
-{% 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;
@@ -48,10 +51,10 @@ fn loadSHMA(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   let col = c_idx * 8u;
   for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
     let k = k_idx + col + col_offset;
-    if (a_global < M) {
 {% if transA %}
       // op(A) = A^T: A stored [.., K, M], so op(A)[a_global, k] = A[k, a_global].
-      tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(a[a_base + k * M + a_global]);
 {% else %}
       tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(a[a_base + a_global * A_M_STRIDE + k]);
 {% endif %}
@@ -62,7 +65,13 @@ fn loadSHMA(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
 }
 fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   let b_col = tile_base + row;
   let col = c_idx * 16u;
   for (var i = 0u; i < 16u; i = i + 1u) {
     let k = k_idx + col + i;
@@ -75,9 +84,19 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   }
 }
-fn storeOutput(offset: u32, row: u32, col: u32, src_slot: u32, row_limit: i32) {
   if (row_limit > 0 && row < u32(row_limit)) {
     let col2 = col + 1u;
     y[offset + row * N + col]       = {{ outScalar }}(ALPHA * scratch[src_slot][0][row * 8u + col]);
     y[offset + row * N + col + 8u]  = {{ outScalar }}(ALPHA * scratch[src_slot][1][row * 8u + col]);
     y[offset + row * N + col + 16u] = {{ outScalar }}(ALPHA * scratch[src_slot][2][row * 8u + col]);
@@ -87,6 +106,7 @@ fn storeOutput(offset: u32, row: u32, col: u32, src_slot: u32, row_limit: i32) {
     y[offset + row * N + col2 + 8u]  = {{ outScalar }}(ALPHA * scratch[src_slot][1][row * 8u + col2]);
     y[offset + row * N + col2 + 16u] = {{ outScalar }}(ALPHA * scratch[src_slot][2][row * 8u + col2]);
     y[offset + row * N + col2 + 24u] = {{ outScalar }}(ALPHA * scratch[src_slot][3][row * 8u + col2]);
   }
 }
@@ -97,6 +117,11 @@ fn main(
   @builtin(subgroup_invocation_id) sg_id: u32,
   @builtin(subgroup_size) sg_size: u32
 ) {
   let batch = workgroup_id.z;
   let a_base = batch * A_BATCH_STRIDE;
   let b_base = batch * B_BATCH_STRIDE;
@@ -125,7 +150,7 @@ fn main(
     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" }});
@@ -160,9 +185,12 @@ fn main(
   workgroupBarrier();
   let row = sg_id / 4u;
   let col = (sg_id % 4u) * 2u;
   var matrix_c_offset = c_base + (a_global_base + base_A) * N + b_global_base + base_B;
   var row_limit = i32(M) - i32(a_global_base + base_A);
-  storeOutput(matrix_c_offset, row, col, subtile_id, row_limit);
   workgroupBarrier();
   subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC10, 8u);
@@ -172,5 +200,5 @@ fn main(
   workgroupBarrier();
   matrix_c_offset = matrix_c_offset + 8u * N;
   row_limit = i32(M) - i32(a_global_base + base_A + 8u);
-  storeOutput(matrix_c_offset, row, col, subtile_id, row_limit);
 }

 // 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.
+// The row count M arrives per call in `params.M`, so one pipeline serves every M;
+// K, N and the batch layout compile in.
+// A K % 32 == 0 gate keeps the reduction loop whole; N is free of the 64-wide column
+// tile because nTailSafe clamps the trailing tile's B columns to N - 1 and guards
+// every store on col < N. Both operands stage through workgroup memory, so
+// subgroupMatrixLoad only ever reads the full tile_A/tile_B arrays and never sees a
+// partial 8x8 tile at any M or N. The clamp is not interchangeable with a zero fill:
+// an out-of-bounds subgroupMatrixLoad resets to offset 0 and returns a different
+// valid tile, and a duplicated real column keeps the discarded accumulators finite.
+// The batch is required to match between A and B (no broadcast) because a_base and
+// 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 %}
 {% set accScalar = "f32" %}
+{% set N_TAIL = nTailSafe is defined and nTailSafe %}
 const K: u32 = {{ K }}u;
 const N: u32 = {{ N }}u;
 {% if transBatchA %}
 const BATCH_COUNT: u32 = {{ batchCount }}u;
 const A_BATCH_STRIDE: u32 = K;
 const A_M_STRIDE: u32 = BATCH_COUNT * K;
+{% elif not transA %}
 const A_M_STRIDE: u32 = K;
 {% endif %}
 const B_BATCH_STRIDE: u32 = K * N;
 const ALPHA: {{ accScalar }} = {{ accScalar }}({{ alpha }});
 const TILE_COLS: u32 = 64u;
 const TILE_ROWS: u32 = 32u;
   let col = c_idx * 8u;
   for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
     let k = k_idx + col + col_offset;
+    if (a_global < params.M) {
 {% if transA %}
       // op(A) = A^T: A stored [.., K, M], so op(A)[a_global, k] = A[k, a_global].
+      tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(a[a_base + k * params.M + a_global]);
 {% else %}
       tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(a[a_base + a_global * A_M_STRIDE + k]);
 {% endif %}
 }
 fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+{% if N_TAIL %}
+  // The trailing column tile reads column N - 1 in place of every column past N.
+  // storeOutput discards those lanes; duplicating a real column keeps them finite.
+  let b_col = min(tile_base + row, N - 1u);
+{% else %}
   let b_col = tile_base + row;
+{% endif %}
   let col = c_idx * 16u;
   for (var i = 0u; i < 16u; i = i + 1u) {
     let k = k_idx + col + i;
   }
 }
+fn storeOutput(offset: u32{% if N_TAIL %}, 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 N_TAIL %}
+{% for block in range(4) %}
+    if (col_base + col + {{ block * 8 }}u < N) {
+      y[offset + row * N + col + {{ block * 8 }}u] = {{ outScalar }}(ALPHA * scratch[src_slot][{{ block }}][row * 8u + col]);
+    }
+    if (col_base + col2 + {{ block * 8 }}u < N) {
+      y[offset + row * N + col2 + {{ block * 8 }}u] = {{ outScalar }}(ALPHA * scratch[src_slot][{{ block }}][row * 8u + col2]);
+    }
+{% endfor %}
+{% else %}
     y[offset + row * N + col]       = {{ outScalar }}(ALPHA * scratch[src_slot][0][row * 8u + col]);
     y[offset + row * N + col + 8u]  = {{ outScalar }}(ALPHA * scratch[src_slot][1][row * 8u + col]);
     y[offset + row * N + col + 16u] = {{ outScalar }}(ALPHA * scratch[src_slot][2][row * 8u + col]);
     y[offset + row * N + col2 + 8u]  = {{ outScalar }}(ALPHA * scratch[src_slot][1][row * 8u + col2]);
     y[offset + row * N + col2 + 16u] = {{ outScalar }}(ALPHA * scratch[src_slot][2][row * 8u + col2]);
     y[offset + row * N + col2 + 24u] = {{ outScalar }}(ALPHA * scratch[src_slot][3][row * 8u + col2]);
+{% endif %}
   }
 }
   @builtin(subgroup_invocation_id) sg_id: u32,
   @builtin(subgroup_size) sg_size: u32
 ) {
+  let M = params.M;
+{% if not transBatchA %}
+  let A_BATCH_STRIDE = M * K;
+{% endif %}
+  let C_BATCH_STRIDE = M * N;
   let batch = workgroup_id.z;
   let a_base = batch * A_BATCH_STRIDE;
   let b_base = batch * B_BATCH_STRIDE;
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set directInputs = false %}
       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" }});
   workgroupBarrier();
   let row = sg_id / 4u;
   let col = (sg_id % 4u) * 2u;
+{% if N_TAIL %}
+  let col_base = b_global_base + base_B;
+{% endif %}
   var matrix_c_offset = c_base + (a_global_base + base_A) * N + b_global_base + base_B;
   var row_limit = i32(M) - i32(a_global_base + base_A);
+  storeOutput(matrix_c_offset{% if N_TAIL %}, col_base{% endif %}, row, col, subtile_id, row_limit);
   workgroupBarrier();
   subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC10, 8u);
   workgroupBarrier();
   matrix_c_offset = matrix_c_offset + 8u * N;
   row_limit = i32(M) - i32(a_global_base + base_A + 8u);
+  storeOutput(matrix_c_offset{% if N_TAIL %}, col_base{% endif %}, row, col, subtile_id, row_limit);
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -20,8 +20,10 @@
   "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 },
@@ -33,46 +35,176 @@
     "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",
     "logicalBShape": "moveAxis(batchMovedBShape, -1, -2) if attrs.transB != 0 and ranks.B > 1 else batchMovedBShape",
     "transBatchContract": "(attrs.transBatchA == 0 and attrs.transBatchB == 0) or (ranks.A == ranks.B and ranks.A >= 3)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
     "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
     "wave32Effective": "wave32Adapter or pinSubgroupSize32",
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
-    "fusedSgmatRank2Ok": "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)",
     "subgroupMatrixResourcesFit": "128 <= deviceWorkgroupCap and ((32 * 32 + 64 * 32) * dtypeBytes(dtypes.T) + 4 * 4 * 64 * 4) <= device.limits.maxComputeWorkgroupStorageSize",
     "sgmatSplitKDepth": "dim(shapes.A, ranks.A - 1)",
-    "sgmatOutTiles": "ceilDiv(dim(shapes.A, 0), 32) * ceilDiv(dim(shapes.B, 1), 64) if fusedSgmatRank2Ok else 1",
-    "sgmatSplitKWant": "ceilDiv(tunables.SUBGROUP_MATRIX_SPLITK_TARGET_WGS, sgmatOutTiles)",
     "sgmatSplitK32Ok": "sgmatSplitKDepth % 1024 == 0",
     "sgmatSplitK16Ok": "sgmatSplitKDepth % 512 == 0",
     "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,
@@ -83,14 +215,11 @@
       },
       "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))"
@@ -101,14 +230,12 @@
           "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" }
         }
       ]
     },
@@ -122,14 +249,11 @@
       },
       "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))"
@@ -140,31 +264,34 @@
           "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": [
@@ -173,21 +300,17 @@
               ]
             }
           ],
-          "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)",
@@ -209,7 +332,7 @@
           "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)",
@@ -222,18 +345,13 @@
     {
       "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": [
         {
@@ -243,24 +361,19 @@
           "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": [
         {
@@ -270,8 +383,7 @@
           "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",
@@ -285,16 +397,12 @@
         ]
       },
       "derive": {
-        "usesF16": "dtypes.T == \"f16\"",
-        "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
-        "scalar": "dtypes.T",
         "hasBias": false,
         "generalAddressing": true,
         "tailSafe": false,
         "outputBuffer": "\"partials\"",
         "outScalar": "\"f32\"",
-        "alpha": "attrs.alpha",
-        "M": "dim(shapes.A, 0)",
         "K": "dim(shapes.A, 1)",
         "N": "dim(shapes.B, 1)",
         "batchCount": 1,
@@ -309,20 +417,15 @@
           "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)",
@@ -342,34 +445,35 @@
       },
       "derive": {
         "hasBias": false,
-        "usesF16": true,
         "fScalar": "\"f16\"",
         "outScalar": "\"f16\"",
-        "scalar": "dtypes.T",
         "generalAddressing": true,
         "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 1) % 64 != 0",
         "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.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": [
@@ -378,14 +482,12 @@
         ]
       },
       "derive": {
-        "usesF16": "dtypes.T == \"f16\"",
-        "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
         "outScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
-        "scalar": "dtypes.T",
         "transA": "attrs.transA != 0",
         "transB": "attrs.transB != 0",
         "transBatchA": "attrs.transBatchA != 0",
-        "M": "(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)))",
         "K": "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1))",
         "N": "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))",
         "batchCount": "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)))"
@@ -395,85 +497,80 @@
           "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)",
-            "z": "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1))"
           }
         }
-      ]
     },
     {
       "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)",
             "z": "dim(shapes.Y, 0)"
           }
         }
@@ -483,7 +580,6 @@
       "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",
@@ -494,16 +590,13 @@
             "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)",
             "z": "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)))))"
           }
         }
@@ -512,27 +605,25 @@
     {
       "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)"
           }
         }

   "typeConstraints": { "T": ["float32", "float16"] },
   "tunables": {
     "TILED_REG_MIN_WORKGROUPS": { "default": 64 },
+    "PLAIN_RANK2_REG_DEEP_K_TILES": { "default": 128 },
     "GEMV_TARGET_BLOCKS": { "default": 512 },
     "SUBGROUP_MATRIX_MIN_M": { "default": 2 },
+    "SUBGROUP_MATRIX_MAX_N_PADDING_RATIO": { "default": 1 },
     "SUBGROUP_MATRIX_SPLITK_TARGET_WGS": { "default": 512 },
     "SUBGROUP_MATRIX_SPLITK_MIN_K": { "default": 1024 },
     "SUBGROUP_MATRIX_SPLITK_MAX_TILES": { "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 },
+    "BROADCAST_TRANSB_MIN_WORKGROUPS": { "default": 64 },
+    "BROADCAST_TRANSB_MAX_PADDING_RATIO": { "default": 2 }
   },
   "derive": {
     "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",
     "logicalBShape": "moveAxis(batchMovedBShape, -1, -2) if attrs.transB != 0 and ranks.B > 1 else batchMovedBShape",
     "transBatchContract": "(attrs.transBatchA == 0 and attrs.transBatchB == 0) or (ranks.A == ranks.B and ranks.A >= 3)",
+    "gemvN": "dim(shapes.B, 1)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
     "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
     "wave32Effective": "wave32Adapter or pinSubgroupSize32",
+    "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32",
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "rank2DeepPortableTier": "variableSubgroup16To32 or (has(device.adapterInfo, \"architecture\") and (device.adapterInfo.architecture == \"pascal\" or (not device.features.has(\"subgroups\") and (device.adapterInfo.architecture == \"apple\" or device.adapterInfo.architecture == \"gen-9\"))))",
+    "broadcastTransbM": "dim(shapes.A, ranks.A - 2)",
+    "broadcastTransbN": "dim(shapes.B, ranks.B - 2)",
+    "broadcastTransbK": "dim(shapes.A, ranks.A - 1)",
+    "broadcastTransbBatches": "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 2))",
+    "gemvLanes": 32,
+    "vec4OutputTile": "4 * gemvLanes",
+    "gemvWorkgroups": "ceilDiv(gemvN, vec4OutputTile)",
+    "gemvSliceCap": "min(32, device.limits.maxComputeWorkgroupSizeY, floor(device.limits.maxComputeInvocationsPerWorkgroup / gemvLanes), floor(device.limits.maxComputeWorkgroupStorageSize / (16 * gemvLanes)))",
+    "gemvSlices": "max(1, min(gemvSliceCap, max(8, pow2ceil(ceilDiv(tunables.GEMV_TARGET_BLOCKS, gemvWorkgroups)))))",
+    "gemvResourcesFit": "gemvLanes <= device.limits.maxComputeWorkgroupSizeX and gemvSlices <= device.limits.maxComputeWorkgroupSizeY and gemvLanes * gemvSlices <= device.limits.maxComputeInvocationsPerWorkgroup and 16 * gemvLanes * gemvSlices <= device.limits.maxComputeWorkgroupStorageSize",
+    "registerTile": 64,
+    "generalTile": 32,
+    "tiledRegResourcesFit": "registerTile / 4 <= device.limits.maxComputeWorkgroupSizeX and registerTile / 4 <= device.limits.maxComputeWorkgroupSizeY and registerTile * registerTile / 16 <= device.limits.maxComputeInvocationsPerWorkgroup and 32 * registerTile * dtypeBytes(dtypes.T) <= device.limits.maxComputeWorkgroupStorageSize",
+    "plainRank2RegDeepPreferredTier": "rank2DeepPortableTier and dim(shapes.A, 1) >= tunables.PLAIN_RANK2_REG_DEEP_K_TILES * 16 and dim(shapes.A, 1) % 16 == 0",
     "subgroupMatrixResourcesFit": "128 <= deviceWorkgroupCap and ((32 * 32 + 64 * 32) * dtypeBytes(dtypes.T) + 4 * 4 * 64 * 4) <= device.limits.maxComputeWorkgroupStorageSize",
     "sgmatSplitKDepth": "dim(shapes.A, ranks.A - 1)",
     "sgmatSplitK32Ok": "sgmatSplitKDepth % 1024 == 0",
     "sgmatSplitK16Ok": "sgmatSplitKDepth % 512 == 0",
     "sgmatSplitK8Ok": "sgmatSplitKDepth % 256 == 0",
     "sgmatSplitK4Ok": "sgmatSplitKDepth % 128 == 0",
     "sgmatSplitK2Ok": "sgmatSplitKDepth % 64 == 0",
+    "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)))",
+    "scalar": "dtypes.T",
+    "alpha": "attrs.alpha",
+    "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+    "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
+    "aRank": "ranks.A",
+    "bRank": "ranks.B",
+    "fusedSgmatRank2Ok": "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)",
+    "fusedSgmatN": "dim(shapes.B, ranks.B - 2) if (attrs.transB != 0 and ranks.B >= 2) else dim(shapes.B, ranks.B - 1)",
+    "fusedSgmatNTail": "fusedSgmatN % 64 != 0",
+    "fusedSgmatNPaddingHigh": "ceilDiv(fusedSgmatN, 64) * 64 > tunables.SUBGROUP_MATRIX_MAX_N_PADDING_RATIO * fusedSgmatN",
+    "sgmatOutTiles": "ceilDiv(dim(shapes.A, 0), 32) * ceilDiv(dim(shapes.B, 1), 64) if fusedSgmatRank2Ok else 1",
+    "sgmatSplitKWant": "ceilDiv(tunables.SUBGROUP_MATRIX_SPLITK_TARGET_WGS, sgmatOutTiles)",
     "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)",
+    "broadcastTransbContract": "(f16Ok(dtypes.T)) and (attrs.transA == 0 and attrs.transB != 0 and attrs.transBatchA == 0 and attrs.transBatchB == 0) and (ranks.A > ranks.B and ranks.B >= 2) and (ranks.Y == ranks.A) and (sameShape(shapes.Y, matmulShape(logicalAShape, logicalBShape))) and (dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 1)) and (dim(shapes.A, ranks.A - 2) >= 64) and (dim(shapes.A, ranks.A - 1) >= 32) and (dim(shapes.B, ranks.B - 2) >= 64) and (numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 2)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535))"
   },
   "bindings": {
+    "a": { "arg": "A", "elementType": "$scalar" },
+    "b": { "arg": "B", "elementType": "$vectorScalar" },
     "partials": { "buffer": "read-only-storage", "elementType": "f32" },
+    "y": { "arg": "Y", "elementType": "$scalar" },
+    "params": { "struct": [{ "name": "cols", "type": "u32", "value": "numel(shapes.Y)" }] },
+    "b_scalar": { "arg": "B", "name": "b", "elementType": "$scalar" },
+    "params_rows": { "name": "params", "struct": [{ "name": "M", "type": "u32", "value": "rowCount" }] }
   },
   "variants": [
+    {
+      "id": "broadcast_transb_tiled_reg",
+      "priority": 6,
+      "when": ["broadcastTransbContract", "tiledRegResourcesFit", "ceilDiv(dim(shapes.A, ranks.A - 2), registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.B, ranks.B - 2), registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": {
+        "bShape": "logicalBShape",
+        "bTransposed": true,
+        "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else broadcastTransbM",
+        "N": "broadcastTransbN",
+        "batchCount": "broadcastTransbBatches if ranks.B > 2 else 1",
+        "transBatchA": false,
+        "regSequentialK": "dtypes.T == \"f16\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.BroadcastTransBTiledReg",
+          "shader": "matmul-tiled-general-reg.wgsl.jinja",
+          "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []",
+            "aRank": "ranks.A if ranks.B > 2 else 2",
+            "K": "dim(shapes.A, ranks.A - 1)"
+          },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceilDiv(N, registerTile)", "y": "ceilDiv(rowCount, registerTile)", "z": "batchCount" }
+        }
+      ],
+      "demoteWhen": ["broadcastTransbBatches * ceilDiv(broadcastTransbM, registerTile) * ceilDiv(broadcastTransbN, registerTile) < tunables.BROADCAST_TRANSB_MIN_WORKGROUPS", "ceilDiv(broadcastTransbM, registerTile) * registerTile * ceilDiv(broadcastTransbN, registerTile) * registerTile * ceilDiv(broadcastTransbK,16) * 16 > tunables.BROADCAST_TRANSB_MAX_PADDING_RATIO * broadcastTransbM * broadcastTransbN * broadcastTransbK"]
+    },
+    {
+      "id": "broadcast_transb_subgroup_matrix_f16",
+      "priority": 11,
+      "when": ["broadcastTransbContract", "dtypes.T == \"f16\"", "wave32Effective", "subgroupMatrixResourcesFit", "ceilDiv(dim(shapes.A, ranks.A - 2),32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.B, ranks.B - 2),64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "bShape": "logicalBShape",
+        "bTransposed": true,
+        "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else broadcastTransbM",
+        "K": "broadcastTransbK",
+        "N": "broadcastTransbN",
+        "batchCount": "broadcastTransbBatches if ranks.B > 2 else 1",
+        "hasBias": false,
+        "fScalar": "dtypes.T",
+        "outScalar": "dtypes.T",
+        "generalAddressing": true,
+        "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 2) % 64 != 0",
+        "outputBuffer": "\"y\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.BroadcastTransBSubgroupMatrix",
+          "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []",
+            "aRank": "ranks.A if ranks.B > 2 else 2"
+          },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceilDiv(N,64)", "y": "ceilDiv(rowCount,32)", "z": "batchCount" }
+        }
+      ],
+      "demoteWhen": ["broadcastTransbBatches * ceilDiv(broadcastTransbM,32) * ceilDiv(broadcastTransbN,64) < tunables.BROADCAST_TRANSB_MIN_WORKGROUPS", "ceilDiv(broadcastTransbM,32) * 32 * ceilDiv(broadcastTransbN,64) * 64 * ceilDiv(broadcastTransbK,32) * 32 > tunables.BROADCAST_TRANSB_MAX_PADDING_RATIO * broadcastTransbM * broadcastTransbN * broadcastTransbK"]
+    },
+    {
+      "id": "broadcast_transb_subgroup_matrix_f32",
+      "priority": 11,
+      "when": ["broadcastTransbContract", "dtypes.T == \"f32\"", "wave32Effective", "subgroupMatrixResourcesFit", "ceilDiv(dim(shapes.A, ranks.A - 2),32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.B, ranks.B - 2),64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "bShape": "logicalBShape",
+        "bTransposed": true,
+        "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else broadcastTransbM",
+        "K": "broadcastTransbK",
+        "N": "broadcastTransbN",
+        "batchCount": "broadcastTransbBatches if ranks.B > 2 else 1",
+        "hasBias": false,
+        "fScalar": "dtypes.T",
+        "outScalar": "dtypes.T",
+        "generalAddressing": true,
+        "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 2) % 64 != 0",
+        "outputBuffer": "\"y\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.BroadcastTransBSubgroupMatrix",
+          "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []",
+            "aRank": "ranks.A if ranks.B > 2 else 2"
+          },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceilDiv(N,64)", "y": "ceilDiv(rowCount,32)", "z": "batchCount" }
+        }
+      ],
+      "demoteWhen": ["broadcastTransbBatches * ceilDiv(broadcastTransbM,32) * ceilDiv(broadcastTransbN,64) < tunables.BROADCAST_TRANSB_MIN_WORKGROUPS", "ceilDiv(broadcastTransbM,32) * 32 * ceilDiv(broadcastTransbN,64) * 64 * ceilDiv(broadcastTransbK,32) * 32 > tunables.BROADCAST_TRANSB_MAX_PADDING_RATIO * broadcastTransbM * broadcastTransbN * broadcastTransbK"]
+    },
     {
       "id": "subgroup_matrix_transbatch_b_f16",
       "priority": 11,
       },
       "derive": {
         "hasBias": false,
         "fScalar": "dtypes.T",
         "outScalar": "dtypes.T",
         "generalAddressing": true,
         "outputBuffer": "\"y\"",
+        "rowCount": "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))"
           "name": "FusedMatMul.SubgroupMatrixTransBatchB",
           "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
           "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2)",
             "bShape": "logicalBShape",
             "bStorageStrides": "[dim(shapes.B, 2), dim(shapes.B, 1) * dim(shapes.B, 2), 1]"
           },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceil(N / 64)", "y": "ceil(rowCount / 32)", "z": "batchCount" }
         }
       ]
     },
       },
       "derive": {
         "hasBias": false,
         "fScalar": "dtypes.T",
         "outScalar": "dtypes.T",
         "generalAddressing": true,
         "outputBuffer": "\"y\"",
+        "rowCount": "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))"
           "name": "FusedMatMul.SubgroupMatrixTransBatchB",
           "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
           "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2)",
             "bShape": "logicalBShape",
             "bStorageStrides": "[dim(shapes.B, 2), dim(shapes.B, 1) * dim(shapes.B, 2), 1]"
           },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceil(N / 64)", "y": "ceil(rowCount / 32)", "z": "batchCount" }
         }
       ]
     },
     {
+      "id": "m1_gemv_vec4",
       "priority": 30,
+      "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "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", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "f16Ok(dtypes.T)"],
+      "derive": {
+        "unrollK2": "dtypes.T == \"f16\"",
+        "gemvScalar": "dtypes.T",
+        "gemvVector": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "alphaScale": "attrs.alpha"
+      },
       "passes": [
         {
           "id": "main",
+          "name": "FusedMatMul.M1GemvVec4",
           "shader": "matmul-vector-matrix-vec4.wgsl.jinja",
           "bindings": [
+            { "arg": "A", "name": "a", "elementType": "$gemvScalar" },
+            { "arg": "B", "name": "b", "elementType": "$gemvVector" },
+            { "arg": "Y", "name": "c", "elementType": "$gemvVector" },
             {
               "name": "params",
               "struct": [
               ]
             }
           ],
+          "dispatch": { "x": "gemvWorkgroups" }
         }
       ]
     },
     {
       "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)", "gemvLanes <= 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", "gemvLanes * tunables.BAND_SPLIT_SLICES <= device.limits.maxComputeInvocationsPerWorkgroup"],
       "derive": {
         "batched": false,
         "outputBuffer": "\"y\"",
         "M": "dim(shapes.A, 0)",
         "K": "dim(shapes.A, 1)",
         "N": "dim(shapes.B, 1)",
           "id": "combine",
           "name": "FusedMatMul.Rank2BandVec4SplitKCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"" },
           "bindings": ["partials", "y", "params"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.Y)), (256)), 65535)",
     {
       "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)", "gemvResourcesFit", "not (gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS and bandSplitK >= 2)"],
       "derive": {
         "batched": false,
         "outputBuffer": "\"y\"",
         "M": "dim(shapes.A, 0)",
         "K": "dim(shapes.A, 1)",
+        "N": "dim(shapes.B, 1)"
       },
       "passes": [
         {
           "bindings": ["a", "b", { "arg": "Y", "name": "y", "elementType": "$vectorScalar" }],
           "dispatch": { "x": "gemvWorkgroups" }
         }
+      ]
     },
     {
       "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)", "gemvResourcesFit", "not (gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS and bandSplitK >= 2)"],
+      "demoteWhen": ["dtypes.T != \"f32\" or (dim(shapes.A, 0) > tunables.BAND_PREFER_MAX_ROWS and dim(shapes.A, 1) >= tunables.BAND_PREFER_DEEP_K)"],
       "derive": {
         "batched": false,
         "outputBuffer": "\"y\"",
         "M": "dim(shapes.A, 0)",
         "K": "dim(shapes.A, 1)",
+        "N": "dim(shapes.B, 1)"
       },
       "passes": [
         {
           "bindings": ["a", "b", { "arg": "Y", "name": "y", "elementType": "$vectorScalar" }],
           "dispatch": { "x": "gemvWorkgroups" }
         }
+      ]
     },
     {
       "id": "subgroup_matrix_splitk",
         ]
       },
       "derive": {
         "hasBias": false,
         "generalAddressing": true,
         "tailSafe": false,
         "outputBuffer": "\"partials\"",
         "outScalar": "\"f32\"",
+        "rowCount": "dim(shapes.A, 0)",
         "K": "dim(shapes.A, 1)",
         "N": "dim(shapes.B, 1)",
         "batchCount": 1,
           "id": "partial",
           "name": "FusedMatMul.SubgroupMatrixSplitK",
           "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": { "bShape": ["dim(shapes.B, 0)", "dim(shapes.B, 1)"], "aRank": 2, "bRank": 2 },
+          "bindings": ["a", "b_scalar", { "name": "partials", "elementType": "f32" }, "params_rows"],
           "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\"" },
           "bindings": ["partials", "y", "params"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.Y)), (256)), 65535)",
       },
       "derive": {
         "hasBias": false,
         "fScalar": "\"f16\"",
         "outScalar": "\"f16\"",
         "generalAddressing": true,
         "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 1) % 64 != 0",
         "outputBuffer": "\"y\"",
+        "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else 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)) if ranks.B > 2 else 1"
       },
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.SubgroupMatrixTailBroadcast",
           "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []",
+            "aRank": "ranks.A if ranks.B > 2 else 2",
+            "bShape": "shapes.B"
+          },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceil(N / 64)", "y": "ceil(rowCount / 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", "(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": [
         ]
       },
       "derive": {
         "outScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
+        "nTailSafe": "fusedSgmatNTail",
         "transA": "attrs.transA != 0",
         "transB": "attrs.transB != 0",
         "transBatchA": "attrs.transBatchA != 0",
+        "rowCount": "(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)))",
         "K": "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1))",
         "N": "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))",
         "batchCount": "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)))"
           "id": "main",
           "name": "FusedMatMul.SubgroupMatrix",
           "shader": "fused-matmul-subgroup-matrix.wgsl.jinja",
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceil(N / 64)", "y": "ceil(rowCount / 32)", "z": "numel(shapes.Y) / (rowCount * N)" }
         }
+      ],
+      "demoteWhen": ["fusedSgmatNPaddingHigh"]
     },
     {
       "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) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, 2) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.BroadcastRank4TiledReg",
           "shader": "matmul-tiled-general-reg.wgsl.jinja",
           "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []",
+            "aRank": "ranks.A if ranks.B > 2 else 2",
+            "K": "dim(shapes.A, ranks.A - 1)",
             "bShape": "shapes.B",
             "transBatchA": "false"
           },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
           "dispatch": {
+            "x": "ceil(dim(shapes.B, ranks.B - 1) / registerTile)",
+            "y": "ceil(rowCount / registerTile)",
+            "z": "batchCount"
           }
         }
+      ],
+      "derive": {
+        "rowCount": "outer(shapes.A, 3) if ranks.B == 2 else dim(shapes.A, 2)",
+        "batchCount": "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) if ranks.B > 2 else 1"
+      }
     },
     {
       "id": "plain_rank2_tiled_reg",
       "priority": 4,
+      "demoteWhen": ["plainRank2RegDeepPreferredTier"],
+      "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) / registerTile) * ceil(dim(shapes.B, 1) / registerTile) >= tunables.TILED_REG_MIN_WORKGROUPS", "ceil(dim(shapes.B, 1) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, 0) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.PlainRank2TiledReg",
           "shader": "matmul-tiled-general-reg.wgsl.jinja",
           "derive": {
+            "rowCount": "dim(shapes.A, 0)",
+            "K": "dim(shapes.A, 1)",
             "bShape": "shapes.B",
             "transBatchA": "false"
           },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": {
+            "x": "ceil(dim(shapes.B, 1) / registerTile)",
+            "y": "ceil(dim(shapes.A, 0) / registerTile)",
+            "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) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, 0) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dim(shapes.Y, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.TransBatchATiledReg",
           "shader": "matmul-tiled-general-reg.wgsl.jinja",
+          "derive": { "aShape": "shapes.A", "bShape": "shapes.B", "transBatchA": "true", "kTile": "4" },
+          "bindings": ["a", "b_scalar", "y"],
           "dispatch": {
+            "x": "ceil(dim(shapes.B, 2) / registerTile)",
+            "y": "ceil(dim(shapes.A, 0) / registerTile)",
             "z": "dim(shapes.Y, 0)"
           }
         }
       "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)"],
       "passes": [
         {
           "id": "main",
             "bShape": "shapes.B",
             "transA": "attrs.transA != 0",
             "transB": "attrs.transB != 0",
             "transBatchA": "attrs.transBatchA != 0",
             "transBatchB": "attrs.transBatchB != 0"
           },
+          "bindings": ["a", "b_scalar", "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)))) / generalTile)",
+            "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)))) / generalTile)",
             "z": "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)))))"
           }
         }
     {
       "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), registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.A, 1), registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dim(shapes.A, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dim(shapes.A, 0) * ceilDiv(dim(shapes.A, 1), registerTile) * ceilDiv(dim(shapes.B, 2), registerTile) >= tunables.TILED_REG_MIN_WORKGROUPS", "tiledRegResourcesFit"],
       "passes": [
         {
           "id": "main",
           "name": "FusedMatMul.TransBatchBTiledReg",
           "shader": "matmul-tiled-general-reg.wgsl.jinja",
           "derive": {
+            "rowCount": "dim(shapes.A, 1)",
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2)",
+            "K": "dim(shapes.A, ranks.A - 1)",
             "bShape": "logicalBShape",
             "transBatchA": false,
             "bStorageStrides": "[dim(shapes.B, 2), dim(shapes.B, 1) * dim(shapes.B, 2), 1]",
             "regSequentialK": "dtypes.T == \"f16\""
           },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
           "dispatch": {
+            "x": "ceilDiv(dim(shapes.B, 2), registerTile)",
+            "y": "ceilDiv(dim(shapes.A, 1), registerTile)",
             "z": "dim(shapes.Y, 0)"
           }
         }

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

@@ -4,15 +4,11 @@
 //
 // 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 %}
@@ -26,7 +22,7 @@ const K_PER_SPLIT: u32 = (K + K_SPLITS - 1u) / K_SPLITS;
 // 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>

 //
 // 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.
 {{ env.wgsl.resourceDeclarations }}
 const K: u32 = {{ K }}u;
 const N4: u32 = {{ N }}u / 4u;
+const LANES: u32 = {{ gemvLanes }}u;
 // 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 %}
 // footprint does not grow with the band.
 var<workgroup> partials: array<vec4<f32>, LANES * SLICES>;
+@compute @workgroup_size({{ gemvLanes }}, {{ gemvSlices }}, 1)
 fn main(
   @builtin(workgroup_id) workgroup_id: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>

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

@@ -8,28 +8,33 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = fScalar %}
 {% set accScalar = "f32" %}
-{% set GENERAL = generalAddressing is defined and generalAddressing %}
 {% set TAIL = tailSafe is defined and tailSafe %}
 {% 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 N: u32 = {{ N }}u;
 const BATCH_COUNT: u32 = {{ batchCount if batchCount is defined else 1 }}u;
@@ -44,7 +49,9 @@ const B_N_STRIDE: u32 = {{ bNStride }}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;
 const TILE_K: u32 = 32u;
@@ -57,10 +64,13 @@ var<workgroup> scratch: array<array<array<{{ accScalar }}, 64>, 4>, 4>;
 fn loadSHMA(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   let a_global = tile_base + row;
   let col = c_idx * 8u;
   for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
     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 %}
@@ -79,10 +89,11 @@ fn loadSHMA(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
 fn loadSHMAKTail(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   let a_global = tile_base + row;
   let col = c_idx * 8u;
   for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
     let k = k_idx + col + col_offset;
-    if (a_global < M && k < K) {
       tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(a[a_base + a_global * A_M_STRIDE + k * A_K_STRIDE]);
     } else {
       tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(0);
@@ -157,6 +168,11 @@ fn main(
   @builtin(subgroup_invocation_id) sg_id: u32,
   @builtin(subgroup_size) sg_size: u32
 ) {
   let b_global_base = workgroup_id.x * TILE_COLS;
   let subtile_id = local_idx / sg_size;
@@ -191,7 +207,7 @@ fn main(
 {% 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 %}
@@ -205,18 +221,18 @@ fn main(
 {% 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;
-{% if aStride.v != 0 %}    a_base = a_base + c{{ axis }} * {{ aStride.v }}u;
 {% endif %}
 {% if bStride.v != 0 %}    b_base = b_base + c{{ axis }} * {{ bStride.v }}u;
 {% endif %}
@@ -249,7 +265,7 @@ fn main(
       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" }});
@@ -278,15 +294,15 @@ fn main(
       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);
@@ -333,8 +349,9 @@ fn main(
     workgroupBarrier();
   }
   // Re-stage workgroup tiles/scratch before the next M-tile iteration reuses them.
-  // The loop bound is workgroup-uniform (M is a compile-time const, num_wg.y and
-  // workgroup_id.y are uniform), so every invocation reaches this barrier together.
   workgroupBarrier();
   }
 }

 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = fScalar %}
 {% set accScalar = "f32" %}
+{% set GENERAL = true %}
 {% set TAIL = tailSafe is defined and tailSafe %}
 {% 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 STATIC_M = M is defined %}
+{% set ROWS = "M" if STATIC_M else "params.M" %}
+{% set ROW_TEST = "a_global < M" if STATIC_M else "row_in" %}
+{% set aDims = (aShape | default([])) if STATIC_M else (aBatchShape | default([])) %}
+{% set kPerSplit = kPerSplit | default(0) %}
 {% set aR = aRank %}
 {% set bR = bRank %}
 {% set aBatchLen = aR - 2 %}
 {% set bBatchLen = bR - 2 %}
 {% set batchRank = aBatchLen %}
+{% set aMStride = K %}
 {% set aKStride = 1 %}
+{% set bKStride = 1 if bTransposed is defined and bTransposed else bShape[bR-1] %}
+{% set bNStride = bShape[bR-2] if bTransposed is defined and bTransposed else 1 %}
 {% if bStorageStrides is defined %}{% set bKStride = bStorageStrides[bR-2] %}{% set bNStride = bStorageStrides[bR-1] %}{% endif %}
+{% if STATIC_M %}
 const M: u32 = {{ M }}u;
+{% endif %}
 const K: u32 = {{ K }}u;
 const N: u32 = {{ N }}u;
 const BATCH_COUNT: u32 = {{ batchCount if batchCount is defined else 1 }}u;
 {% if TAIL %}const K_FULL: u32 = (K / 32u) * 32u;
 {% endif %}
 const ALPHA: f32 = f32({{ alpha }});
+{% if STATIC_M %}
 const C_BATCH_STRIDE: u32 = M * N;
+{% endif %}
 const TILE_COLS: u32 = 64u;
 const TILE_ROWS: u32 = 32u;
 const TILE_K: u32 = 32u;
 fn loadSHMA(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   let a_global = tile_base + row;
+{% if not STATIC_M %}
+  let row_in = a_global < params.M;
+{% endif %}
   let col = c_idx * 8u;
   for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
     let k = k_idx + col + col_offset;
+    if ({{ ROW_TEST }}) {
 {% 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 %}
 fn loadSHMAKTail(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   let a_global = tile_base + row;
+  let row_in = a_global < {{ ROWS }};
   let col = c_idx * 8u;
   for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
     let k = k_idx + col + col_offset;
+    if (row_in && k < K) {
       tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(a[a_base + a_global * A_M_STRIDE + k * A_K_STRIDE]);
     } else {
       tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(0);
   @builtin(subgroup_invocation_id) sg_id: u32,
   @builtin(subgroup_size) sg_size: u32
 ) {
+{% if not STATIC_M %}
+  // The row count arrives per call; the strides it scales are formed here.
+  let M = params.M;
+  let C_BATCH_STRIDE = M * N;
+{% endif %}
   let b_global_base = workgroup_id.x * TILE_COLS;
   let subtile_id = local_idx / sg_size;
 {% set axis = batchRank - 1 - i %}
 {% set aAxis = axis - (batchRank - aBatchLen) %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aDim = aDims[aAxis] %}
 {% set bDim = bShape[bAxis] if bAxis >= 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 aDim = aDims[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 - 2) %}{% set aStride.v = aStride.v * aDims[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;
+{% if aStride.v != 0 %}    a_base = a_base + c{{ axis }} * {% if aStride.v != 1 %}{{ aStride.v }}u * {% endif %}M * K;
 {% endif %}
 {% if bStride.v != 0 %}    b_base = b_base + c{{ axis }} * {{ bStride.v }}u;
 {% endif %}
       workgroupBarrier();
       for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set directInputs = false %}
       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" }});
       workgroupBarrier();
       for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set directInputs = false %}
+      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);
     workgroupBarrier();
   }
   // Re-stage workgroup tiles/scratch before the next M-tile iteration reuses them.
+  // The loop bound is workgroup-uniform (M is a compile-time const or a uniform
+  // read, and num_wg.y and workgroup_id.y are uniform), so every invocation
+  // reaches this barrier together.
   workgroupBarrier();
   }
 }

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

@@ -3,28 +3,33 @@
 // Register-blocked MatMul for the no-subgroup-matrix
 // tier: Y = alpha * A @ B. It retains the bounds-checked addressing and
 // batch-broadcast of the general kernel, and its transposed-batch-A layout.
-// Transposed and 1-D operands select other variants and are not handled here.
 // Each thread computes a 4x4 micro-tile within a 64x64 workgroup tile, reusing
 // each staged operand across four accumulators. Both tiles are indexed by their
 // own output axis and group four K values per vector word, so the micro-tile
 // accumulates through dot() and one step reads TM + TN words 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
@@ -39,7 +44,9 @@
 {% 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_M_STRIDE: u32 = {{ aMStride }}u;
@@ -50,11 +57,13 @@ const B_N_STRIDE: u32 = {{ bNStride }}u;
 // 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 BK: u32 = {{ kTile }}u;
-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
 {% if splitK > 1 %}
 const SPLIT_K: u32 = {{ splitK }}u;
 const K_PER_SPLIT: u32 = {{ kPerSplit }}u;
@@ -70,19 +79,22 @@ var<workgroup> tileB: array<array<vec4<{{ tileT }}>, K_VECS>, BN>;   // B[n][k/4
 {% 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 %}
-@compute @workgroup_size(16, 16, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let mBase = wg.y * BM;
   let nBase = wg.x * BN;
-  let li = lid.y * 16u + lid.x;
   let zOut = wg.z;
 {% if splitK > 1 %}
@@ -100,17 +112,17 @@ fn main(
 {% 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;
-{% if aStride.v != 0 %}  aBatchOff = aBatchOff + c{{ axis }} * {{ aStride.v }}u;
 {% endif %}
 {% if bStride.v != 0 %}  bBatchOff = bBatchOff + c{{ axis }} * {{ bStride.v }}u;
 {% endif %}
@@ -134,7 +146,7 @@ fn main(
 {% 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) {
       let ar = idx / K_VECS;
       let ac4 = idx % K_VECS;
       let am = mBase + ar;
@@ -148,7 +160,7 @@ fn main(
       }
       tileA[ar][ac4] = aWord;
     }
-    for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + 256u) {
       let bc = idx % BN;
       let br4 = idx / BN;
       let bn = nBase + bc;
@@ -164,6 +176,7 @@ fn main(
     }
     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) {
@@ -173,9 +186,12 @@ for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
   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 %}
     }
   }

 // Register-blocked MatMul for the no-subgroup-matrix
 // tier: Y = alpha * A @ B. It retains the bounds-checked addressing and
 // batch-broadcast of the general kernel, and its transposed-batch-A layout.
+// A retains its matrix-axis order; B may transpose its matrix axes.
+// 1-D operands use other variants.
 // Each thread computes a 4x4 micro-tile within a 64x64 workgroup tile, reusing
 // each staged operand across four accumulators. Both tiles are indexed by their
 // own output axis and group four K values per vector word, so the micro-tile
 // accumulates through dot() and one step reads TM + TN words rather than
+// 4 * (TM + TN) scalars. Logical matrix axes specialize to physical operand strides.
 {% set aR = aRank %}
 {% set bR = bRank %}
 {% set aBatchLen = aR - 2 %}
 {% set bBatchLen = bR - 2 %}
 {% set batchRank = aBatchLen %}
+{% set STATIC_M = aShape is defined %}
+{% set aDims = aShape if aShape is defined else (aBatchShape | default([])) %}
 {% set aTailStride = namespace(v=1) %}
+{% if aShape is defined %}{% for j in range(1, aR) %}{% set aTailStride.v = aTailStride.v * aShape[j] %}{% endfor %}{% endif %}
 {% set bTailStride = namespace(v=1) %}
 {% for j in range(1, bR) %}{% set bTailStride.v = bTailStride.v * bShape[j] %}{% endfor %}
+{% if aShape is defined %}
 {% if transBatchA %}{% set M = aShape[0] %}{% set K = aShape[aR-1] %}
 {% else %}{% set M = aShape[aR-2] %}{% set K = aShape[aR-1] %}{% endif %}
+{% endif %}
 {% set N = bShape[bR-1] %}
 {% if transBatchA %}{% set aMStride = aTailStride.v %}{% set aKStride = 1 %}
+{% else %}{% set aMStride = K %}{% set aKStride = 1 %}{% endif %}
+{% set bKStride = 1 if bTransposed is defined and bTransposed else bShape[bR-1] %}
+{% set bNStride = bShape[bR-2] if bTransposed is defined and bTransposed else 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
 {% endif %}
 {% set tileT = scalar if scalar == "f16" else accT %}
 {% set kTile = kTile if kTile is defined else 16 %}
+{% if STATIC_M %}
 const M: u32 = {{ M }}u;
+{% endif %}
 const K: u32 = {{ K }}u;
 const N: u32 = {{ N }}u;
 const A_M_STRIDE: u32 = {{ aMStride }}u;
 // 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.
+{% set microTile = 4 %}
+{% set lanes = 16 %}
 const BK: u32 = {{ kTile }}u;
+const BM: u32 = {{ registerTile }}u;
+const BN: u32 = {{ registerTile }}u;
+const TM: u32 = {{ microTile }}u;   // per-thread micro-tile rows
+const TN: u32 = {{ microTile }}u;   // per-thread micro-tile cols
 {% if splitK > 1 %}
 const SPLIT_K: u32 = {{ splitK }}u;
 const K_PER_SPLIT: u32 = {{ kPerSplit }}u;
 {% set bAxis = axis - (batchRank - bBatchLen) %}
 {% set aStored = (aAxis + 1) if (transBatchA and aAxis >= 0) else aAxis %}
 {% set bStored = bAxis %}
+{% set aDim = aDims[aStored] %}
 {% set bDim = bShape[bStored] if bStored >= 0 else 1 %}
 {% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
 {% endfor %}
+@compute @workgroup_size({{ lanes }}, {{ lanes }}, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
+{% if not STATIC_M %}
+  let M = params.M;
+{% endif %}
   let mBase = wg.y * BM;
   let nBase = wg.x * BN;
+  let li = lid.y * {{ lanes }}u + lid.x;
   let zOut = wg.z;
 {% if splitK > 1 %}
 {% set bAxis = axis - (batchRank - bBatchLen) %}
 {% set aStored = (aAxis + 1) if (transBatchA and aAxis >= 0) else aAxis %}
 {% set bStored = bAxis %}
+{% set aDim = aDims[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 if STATIC_M else aR - 2) %}{% set aStride.v = aStride.v * aDims[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;
+{% if aStride.v != 0 %}  aBatchOff = aBatchOff + c{{ axis }} * {% if STATIC_M %}{{ aStride.v }}u{% else %}{% if aStride.v != 1 %}{{ aStride.v }}u * {% endif %}M * K{% endif %};
 {% endif %}
 {% if bStride.v != 0 %}  bBatchOff = bBatchOff + c{{ axis }} * {{ bStride.v }}u;
 {% endif %}
 {% 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 + {{ lanes * lanes }}u) {
       let ar = idx / K_VECS;
       let ac4 = idx % K_VECS;
       let am = mBase + ar;
       }
       tileA[ar][ac4] = aWord;
     }
+    for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + {{ lanes * lanes }}u) {
       let bc = idx % BN;
       let br4 = idx / BN;
       let bn = nBase + bc;
     }
     workgroupBarrier();
 {% set regT = accT %}{% filter indent(4, true) %}
+{% set regAccumulator = "acc" %}
 let aRow = lid.y * TM;
 let bCol = lid.x * TN;
 for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
   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) %}
+      {{ regAccumulator }}[i * TN + j] = {{ regAccumulator }}[i * TN + j] + av[i][{{ component }}] * bv[j][{{ component }}];
+{% endfor %}
+{% else %}
+      {{ regAccumulator }}[i * TN + j] = {{ regAccumulator }}[i * TN + j] + dot(av[i], bv[j]);
 {% endif %}
     }
   }

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

@@ -18,9 +18,10 @@
  * [d1, ..., dR-2, d0, dR-1]. Stored axis zero becomes the M axis, the final K
  * 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] %}
@@ -34,8 +35,8 @@
 {% 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 %}
@@ -58,12 +59,13 @@ const ALPHA: {{ accT }} = {{ accT }}(1);{% else %}const ALPHA: f32 = f32({{ alph
 // 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.
-const BK: u32 = 16u;
-const BM: u32 = 32u;
-const BN: u32 = 32u;
-var<workgroup> tileA: array<array<{{ tileT }}, 16>, 32>;
-var<workgroup> tileB: array<array<{{ tileT }}, 32>, 16>;
 {% set hasBatchCoord = namespace(value=false) %}
 {% for i in range(batchRank) %}
 {% set axis = batchRank - 1 - i %}
@@ -71,19 +73,19 @@ var<workgroup> tileB: array<array<{{ tileT }}, 32>, 16>;
 {% 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 %}
-@compute @workgroup_size(16, 16, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let mBase = wg.y * BM;
   let nBase = wg.x * BN;
-  let li = lid.y * 16u + lid.x;
   // Per-batch base offsets into A and B using right-aligned broadcast strides.
   // Decompose the flat output-batch index from the innermost axis outward.
@@ -99,11 +101,11 @@ fn main(
 {% 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 %}
@@ -127,7 +129,7 @@ fn main(
     let kBase = kt * BK;
     // Cooperative load: 32x16 A tile + 16x32 B tile, 256 threads x 2 each.
     for (var e: u32 = 0u; e < 2u; e = e + 1u) {
-      let idx = li + e * 256u;
       let ar = idx / BK;
       let ac = idx % BK;
       let am = mBase + ar;

  * [d1, ..., dR-2, d0, dR-1]. Stored axis zero becomes the M axis, the final K
  * axis is unchanged, and the remaining axes form the batch. This stride
  * permutation composes with the ordinary last-two-axis transpose. */
+{% set STATIC_M = true %}
+{% set aDims = aShape if aShape is defined else (aBatchShape | default([])) %}
 {% 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] %}
 {% 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 = M %}
+{% else %}{% set aMStride = K %}{% 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 %}
 // 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 lanes = 16 %}
+const BK: u32 = {{ lanes }}u;
+const BM: u32 = {{ generalTile }}u;
+const BN: u32 = {{ generalTile }}u;
+var<workgroup> tileA: array<array<{{ tileT }}, {{ lanes }}>, {{ generalTile }}>;
+var<workgroup> tileB: array<array<{{ tileT }}, {{ generalTile }}>, {{ lanes }}>;
 {% set hasBatchCoord = namespace(value=false) %}
 {% for i in range(batchRank) %}
 {% set axis = batchRank - 1 - i %}
 {% 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 = aDims[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 %}
+@compute @workgroup_size({{ lanes }}, {{ lanes }}, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let mBase = wg.y * BM;
   let nBase = wg.x * BN;
+  let li = lid.y * {{ lanes }}u + lid.x;
   // Per-batch base offsets into A and B using right-aligned broadcast strides.
   // Decompose the flat output-batch index from the innermost axis outward.
 {% 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 = aDims[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 if STATIC_M else aR - 2) %}{% set aStride.v = aStride.v * aDims[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 %}
     let kBase = kt * BK;
     // Cooperative load: 32x16 A tile + 16x32 B tile, 256 threads x 2 each.
     for (var e: u32 = 0u; e < 2u; e = e + 1u) {
+      let idx = li + e * {{ lanes * lanes }}u;
       let ar = idx / BK;
       let ac = idx % BK;
       let am = mBase + ar;

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

@@ -3,14 +3,14 @@
 {{ 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>;
-@compute @workgroup_size(32, {{ gemvSlices }}, 1)
 fn main(
   @builtin(workgroup_id) workgroup_id: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
@@ -20,9 +20,23 @@ fn main(
   let cg = workgroup_id.x * LANES + lane;
   var acc = vec4<f32>(0.0);
   if (cg < params.N4) {
     for (var k = slice; k < params.K; k = k + SLICES) {
       acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
     }
   }
   partials[slice * LANES + lane] = acc;
   workgroupBarrier();

 {{ env.wgsl.resourceDeclarations }}
 {% set OUT = outputBuffer if outputBuffer is defined else "c" %}
+const LANES: u32 = {{ gemvLanes }}u;
 // 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>;
+@compute @workgroup_size({{ gemvLanes }}, {{ gemvSlices }}, 1)
 fn main(
   @builtin(workgroup_id) workgroup_id: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
   let cg = workgroup_id.x * LANES + lane;
   var acc = vec4<f32>(0.0);
   if (cg < params.N4) {
+{% if unrollK2 is defined and unrollK2 %}
+    // Two K positions per iteration amortize loop/address arithmetic on long
+    // decode projections while preserving each slice's exact strided order.
+    var k = slice;
+    for (; k + SLICES < params.K; k = k + 2u * SLICES) {
+      acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
+      let k1 = k + SLICES;
+      acc = acc + f32(a[k1]) * vec4<f32>(b[k1 * params.N4 + cg]);
+    }
+    if (k < params.K) {
+      acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
+    }
+{% else %}
     for (var k = slice; k < params.K; k = k + SLICES) {
       acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
     }
+{% endif %}
   }
   partials[slice * LANES + lane] = acc;
   workgroupBarrier();

build/webgpu/metadata.json CHANGED Viewed

@@ -1,31 +1,34 @@
 {
   "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"],

 {
   "name": "com.microsoft.FusedMatMul",
+  "id": "_com_microsoft_fusedmatmul_webgpu_17d6f48",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "hPORMosjCjoEdszLuJCkanMmm8j/WxJFR8aS4bZRdWw=",
+      "fused-matmul-subgroup-matrix.wgsl.jinja": "idFWg6H0GgKnO1AG8WFLO2rXeN4cqOwW+V9dJii9Fn4=",
+      "manifest.json": "epnZxQxVo55octrjWqNWM/cs0DsR74Cz59XvUB+MyOI=",
+      "matmul-band-vec4.wgsl.jinja": "LKAs6A++OJF0ZEEM4JITZmKXkr/wo9gd5qrR3zDZy1g=",
+      "matmul-subgroup-matrix-ext.wgsl.jinja": "MGpbY/2ZIE3EEnG0SmKB74zCrPmVGbnNBSw9Ml81+7c=",
+      "matmul-tiled-general-reg.wgsl.jinja": "DGpeWkL4fuu22CI4+Yw3/QD2NON89wCggbotpFNWobE=",
+      "matmul-tiled-general.wgsl.jinja": "j+8KKAFZDVh9M3zNFxJCU4ihsiLgI0KxpE9CD8TqF6Q=",
+      "matmul-vector-matrix-vec4.wgsl.jinja": "Qiv2AO8MWj1BAoPLMCH98oXBasrQYFWFh/IUhdhZp5I=",
+      "reduce-axis0-splitk-combine.wgsl.jinja": "Zf7tz8nrapi2KMZIbv8cT2f4pliJoJHj4iwS4hc/6ms=",
+      "test.json": "LCHLBZTTSQFp/nqb55OJuX7oiL59ThEfHhocijVk5DE="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
+      "broadcast_transb_tiled_reg": ["matmul-tiled-general-reg.wgsl.jinja"],
+      "broadcast_transb_subgroup_matrix_f16": ["matmul-subgroup-matrix-ext.wgsl.jinja"],
+      "broadcast_transb_subgroup_matrix_f32": ["matmul-subgroup-matrix-ext.wgsl.jinja"],
       "subgroup_matrix_transbatch_b_f16": ["matmul-subgroup-matrix-ext.wgsl.jinja"],
       "subgroup_matrix_transbatch_b_f32": ["matmul-subgroup-matrix-ext.wgsl.jinja"],
+      "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"],

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

@@ -2,37 +2,19 @@
 // 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;
@@ -48,7 +30,9 @@ fn is_nan_f32(value: f32) -> bool {
 @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" %}
@@ -74,13 +58,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
     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" %}
@@ -89,7 +66,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
     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];
@@ -101,9 +77,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
       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" %}

 // 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 yv = "f16(" if outputF16 else "" %}
 {% set vy = ")" if outputF16 else "" %}
 {{ 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. */
 fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
+}{% endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
 {% if op == "logsumexp" %}
 const F32_MIN: f32 = -3.4028234663852886e38;
 const F32_MAX: f32 = 3.4028234663852886e38;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
+  // The start already folds gid.y in, so the stride must span every y row too;
+  // an x-only stride would send y = 0 lanes over columns the y >= 1 rows own.
+  let stride = nwg.x * nwg.y * 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" %}
     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 op == "max" %}
     var total = reduction_identity();
 {% elif op == "min" %}
     var total = 1.0;
 {% else %}
     var total = 0.0;
 {% endif %}
     for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
       let p = partials[seg * params.cols + col];
       total = total + p;
 {% endif %}
     }
 {% if op == "l2" %}
     y[col] = {{ yv }}sqrt(total){{ vy }};
 {% elif op == "logsum" %}

build/webgpu/test.json CHANGED Viewed

@@ -7,6 +7,25 @@
     "ort_float32_trans_batch_b_input_B": [1, 0, 1, 2, 0, 1, -1, 0, 1, 1, 0, 1, 2, -1, 1, 1]
   },
   "cases": [
     {
       "name": "ort_float32_broadcast_rank4_by_rank3",
       "provenance": {
@@ -693,7 +712,7 @@
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
         "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
-        "notes": "M=32, K=32, N=64 selects the subgroup-matrix path; finite subnormal dot products must not flush to zero."
       },
       "inputs": {
         "A": { "dtype": "float32", "shape": [32, 32], "data": { "kind": "constant", "value": 1e-20 } },
@@ -936,7 +955,7 @@
     {
       "name": "f32_decode_gemv_m1_k65_n68_vec4_compact",
       "provenance": {
-        "notes": "Compact M=1 float32 GEMV correctness lock for the model-shaped K=4096,N=4096 bandwidth-bound benchmark. Odd K preserves the sliced reduction while N=68 exercises the final partial 128-column workgroup."
       },
       "attrs": { "alpha": 1 },
       "inputs": {
@@ -953,10 +972,50 @@
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [1, 68], "tolerance": 0.0002 } }
     },
     {
       "name": "f32_rank4_by_rank2_shared_weight_compact",
       "provenance": {
-        "notes": "Compact rank-4 by rank-2 shared-weight broadcast lock for the attention-shaped benchmark. Odd M/K/N exercise batch offset and tile-tail handling."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
@@ -975,7 +1034,9 @@
     },
     {
       "name": "subgroup_matrix_kn_tail_f16_compact",
-      "provenance": { "notes": "Compact f16 subgroup-matrix lock with both a partial K=34 tile and N=66 output tail." },
       "attrs": { "alpha": 0.5 },
       "inputs": {
         "A": {
@@ -994,7 +1055,7 @@
     {
       "name": "subgroup_matrix_broadcast_rank4x3_f16_compact",
       "provenance": {
-        "notes": "Compact rank-4 by rank-3 broadcast lock for the model-shaped [1,8,M,K] x [8,K,N] stress case."
       },
       "attrs": { "alpha": 1 },
       "inputs": {
@@ -1014,7 +1075,7 @@
     {
       "name": "broadcast_rank4_tiled_reg_f16_compact",
       "provenance": {
-        "notes": "Compact rank-4 by rank-3 broadcast lock for the register-blocked non-subgroup-matrix path. Odd M/K/N exercise every output and reduction tail."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
@@ -1054,7 +1115,7 @@
     {
       "name": "transbatch_a_dense_m_tail_f16_compact",
       "provenance": {
-        "notes": "Compact lock for stored [M,batch,K] transBatchA addressing. M=65 exercises the subgroup-matrix row tail; all/no-mma/no-subgroups select the MMA/register-blocked portable paths used by the model-shaped stress case."
       },
       "attrs": { "alpha": 0.5, "transBatchA": 1 },
       "inputs": {
@@ -1109,7 +1170,7 @@
       "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 2, 2, 4], "tolerance": 0.000001 } }
     },
     {
-      "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."
       },
@@ -1129,7 +1190,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [33, 66], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
-      "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."
       },
@@ -1169,7 +1230,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 33, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
-      "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."
       },
@@ -1189,7 +1250,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 65, 67], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
-      "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."
       },
@@ -1209,7 +1270,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [2, 65, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
-      "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."
       },
@@ -1229,7 +1290,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
-      "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."
       },
@@ -1248,7 +1309,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [3, 7], "tolerance": 0.01, "relTolerance": 0.005 } }
     },
     {
-      "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."
       },
@@ -1267,7 +1328,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
-      "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."
       },
@@ -1286,7 +1347,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [2, 64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
-      "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."
       },
@@ -1305,7 +1366,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [57, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
-      "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."
       },
@@ -1325,7 +1386,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
-      "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."
       },
@@ -1345,7 +1406,7 @@
       "outputs": { "Y": { "dtype": "float16", "shape": [50, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
-      "name": "f16_rank3_by_broadcast_rank3_offset_scale_lock",
       "provenance": {
         "notes": "f16_rank3_by_broadcast_rank3 cancels to 0.076 under a 0.02 absolute tolerance (26% blind). Offsetting both operands makes each element ~K * aOffset * bOffset over K=3, so the shared single-batch B - read by both output batches - is pinned for value as well as for broadcast addressing."
       },
@@ -1470,7 +1531,7 @@
     {
       "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": {
@@ -1544,7 +1605,7 @@
     {
       "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": {
@@ -1564,7 +1625,7 @@
     {
       "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": {
@@ -1584,7 +1645,7 @@
     {
       "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": {
@@ -1652,7 +1713,7 @@
       },
       "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."
       }
     },
     {
@@ -1672,7 +1733,7 @@
       },
       "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."
       }
     },
     {
@@ -1692,7 +1753,7 @@
       },
       "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."
       }
     },
     {
@@ -1712,7 +1773,7 @@
       },
       "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."
       }
     },
     {
@@ -1732,7 +1793,7 @@
       },
       "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."
       }
     },
     {
@@ -1752,8 +1813,612 @@
       },
       "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."
       }
     }
   ]
 }

     "ort_float32_trans_batch_b_input_B": [1, 0, 1, 2, 0, 1, -1, 0, 1, 1, 0, 1, 2, -1, 1, 1]
   },
   "cases": [
+    {
+      "name": "f16_plain_register_tile_m512_k64_n512",
+      "provenance": {
+        "notes": "A compact rank-two float16 matrix product with offset operands exposes dropped reduction terms across output tiles."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05, "offset": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [512, 512], "tolerance": 0.01, "relTolerance": 0.001 } }
+    },
     {
       "name": "ort_float32_broadcast_rank4_by_rank3",
       "provenance": {
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
         "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "M=32, K=32, N=64 with finite subnormal dot products that must not flush to zero."
       },
       "inputs": {
         "A": { "dtype": "float32", "shape": [32, 32], "data": { "kind": "constant", "value": 1e-20 } },
     {
       "name": "f32_decode_gemv_m1_k65_n68_vec4_compact",
       "provenance": {
+        "notes": "A compact M=1 float32 matrix product with odd K and N=68 checks the reduction and final output-column tail."
       },
       "attrs": { "alpha": 1 },
       "inputs": {
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [1, 68], "tolerance": 0.0002 } }
     },
+    {
+      "name": "f16_decode_gemv_m1_k65_n68_vec4_compact",
+      "provenance": {
+        "notes": "A float16 M=1 matrix product with odd K and N=68 checks both reduction and output-column tails. The dot product accumulates in float32 and narrows only at the store."
+      },
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [1, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [65, 68],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 68], "tolerance": 0.004, "relTolerance": 0.004 } }
+    },
+    {
+      "name": "f16_decode_gemv_m1_k64_n128_alpha_half",
+      "provenance": {
+        "notes": "A float16 M=1 matrix product with even K and non-unit alpha checks complete reduction. The dot product accumulates in float32 and narrows only at the store."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 128], "tolerance": 0.0002, "relTolerance": 0.004 } }
+    },
     {
       "name": "f32_rank4_by_rank2_shared_weight_compact",
       "provenance": {
+        "notes": "A rank-2 weight shared across a 2x3 batch multiplies M=5, K=7, N=9 (alpha=0.5); the odd dimensions check batch-offset indexing and non-power-of-two tails."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
     },
     {
       "name": "subgroup_matrix_kn_tail_f16_compact",
+      "provenance": {
+        "notes": "M=33 and K=34 are one and two past a 32-element boundary and N=66 is two past a 64-element boundary (float16, alpha=0.5), checking partial tiles in all three dimensions."
+      },
       "attrs": { "alpha": 0.5 },
       "inputs": {
         "A": {
     {
       "name": "subgroup_matrix_broadcast_rank4x3_f16_compact",
       "provenance": {
+        "notes": "A rank-4 [1,2,33,32] by rank-3 [2,32,64] float16 broadcast multiplies M=33, K=32, N=64 across a batch of 2, checking batch-dimension broadcasting at a compact scale."
       },
       "attrs": { "alpha": 1 },
       "inputs": {
     {
       "name": "broadcast_rank4_tiled_reg_f16_compact",
       "provenance": {
+        "notes": "Compact rank-4 by rank-3 float16 broadcast with odd M/K/N checks output and reduction tails."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
     {
       "name": "transbatch_a_dense_m_tail_f16_compact",
       "provenance": {
+        "notes": "A stored [M=65, batch=2, K=32] with transBatchA (float16, alpha=0.5): M=65 is one row past a 64-element boundary (K=32,N=64), checking the partial final row across both batches under this transposed-batch addressing."
       },
       "attrs": { "alpha": 0.5, "transBatchA": 1 },
       "inputs": {
       "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 2, 2, 4], "tolerance": 0.000001 } }
     },
     {
+      "name": "subgroup_matrix_kn_tail_f16_offset_alpha_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [33, 66], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
+      "name": "subgroup_matrix_broadcast_rank4x3_f16_offset_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 33, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
+      "name": "broadcast_rank4_tiled_reg_f16_offset_alpha_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 65, 67], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
+      "name": "transbatch_a_dense_m_tail_f16_offset_alpha_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [2, 65, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
+      "name": "aligned_f16_transA_transB_alpha_offset_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
+      "name": "f16_unaligned_3x5x7_offset_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [3, 7], "tolerance": 0.01, "relTolerance": 0.005 } }
     },
     {
+      "name": "aligned_f16_plain_64x32x64_offset_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
+      "name": "aligned_f16_batched_plain_2x64x32x64_offset_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [2, 64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
+      "name": "subgroup_matrix_m_tail_57_partial_block_f16_offset_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [57, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
+      "name": "aligned_f16_transA_64x32_offset_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
+      "name": "transA_transB_subgroup_matrix_m_tail_50_f16_offset_scale",
       "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."
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [50, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
     },
     {
+      "name": "f16_rank3_by_broadcast_rank3_offset_scale",
       "provenance": {
         "notes": "f16_rank3_by_broadcast_rank3 cancels to 0.076 under a 0.02 absolute tolerance (26% blind). Offsetting both operands makes each element ~K * aOffset * bOffset over K=3, so the shared single-batch B - read by both output batches - is pinned for value as well as for broadcast addressing."
       },
     {
       "name": "band_vec4_alpha_scaled_m8_k256_n512",
       "provenance": {
+        "notes": "Eight rows over 512 output columns with non-unit alpha check independent matrix-product results."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
     {
       "name": "f16_rank4_by_rank2_shared_weight_m33_k34_n66",
       "provenance": {
+        "notes": "A [2, 3, 33, 34] float16 A shares one [34, 66] B across both batch axes, with alpha 0.5; none of M=33, K=34 or N=66 is a multiple of 32."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
     {
       "name": "f16_rank4_by_rank2_shared_weight_m8_k32_n64",
       "provenance": {
+        "notes": "A [2, 3, 8, 32] float16 A shares one [32, 64] B across both batch axes, with alpha 0.5, at an 8-row M."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
     {
       "name": "f16_rank4_by_rank2_shared_weight_m1_k32_n64",
       "provenance": {
+        "notes": "A [2, 3, 1, 32] float16 A shares one [32, 64] B across both batch axes, with alpha 0.5, at a single-row M."
       },
       "attrs": { "alpha": 0.5 },
       "inputs": {
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [4, 128, 512], "tolerance": 0.001, "relTolerance": 0.001 } },
       "provenance": {
+        "notes": "Physical B stored [K=128, batch=4, N=512] with transBatchB (float16, alpha=-0.5, negative to avoid cancellation with positive inputs): M=128, K=128, N=512 are all tile-aligned, checking this transposed-batch addressing at a fully aligned shape."
       }
     },
     {
       },
       "outputs": { "Y": { "dtype": "float16", "shape": [3, 129, 513], "tolerance": 0.001, "relTolerance": 0.001 } },
       "provenance": {
+        "notes": "Physical B stored [K=131, batch=3, N=513] with transBatchB (float16, alpha=-0.5): M=129, K=131, N=513 are each one to three past a 128/512 boundary, checking partial final tiles under transposed-batch addressing."
       }
     },
     {
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [4, 128, 512], "tolerance": 0.0001, "relTolerance": 0.00002 } },
       "provenance": {
+        "notes": "Physical B stored [K=128, batch=4, N=512] with transBatchB (float32, alpha=-0.5, negative to avoid cancellation with positive inputs): M=128, K=128, N=512 are all tile-aligned, checking this transposed-batch addressing at a fully aligned shape."
       }
     },
     {
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [3, 129, 513], "tolerance": 0.0001, "relTolerance": 0.00002 } },
       "provenance": {
+        "notes": "Physical B stored [K=131, batch=3, N=513] with transBatchB (float32, alpha=-0.5): M=129, K=131, N=513 are each one to three past a 128/512 boundary, checking partial final tiles under transposed-batch addressing."
       }
     },
     {
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [4, 4096], "tolerance": 0.0001, "relTolerance": 0.00001 } },
       "provenance": {
+        "notes": "M=4, K=2048, N=4096 float32 operands with alpha=0.5 and different per-operand offsets (0.02 vs 0.03) avoid cancellation in the matrix product."
       }
     },
     {
       },
       "outputs": { "Y": { "dtype": "float32", "shape": [16, 4096], "tolerance": 0.0001, "relTolerance": 0.00001 } },
       "provenance": {
+        "notes": "M=16, K=2560, N=4096 float32 operands with alpha=0.5 and different per-operand offsets (0.02 vs 0.03) avoid cancellation in the matrix product."
       }
+    },
+    {
+      "name": "broadcast_transb_rank4x3_aligned_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 64, 64], "tolerance": 0.0005, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_rank4x3_tails_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 67, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 65, 67], "tolerance": 0.0005, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_broadcast_a_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 1, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 67, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 65, 67], "tolerance": 0.0005, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_broadcast_b_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [1, 67, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 65, 67], "tolerance": 0.0005, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_rank3x2_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [3, 64, 64], "tolerance": 0.0005, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_rank5x3_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 1, 3, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 67, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 1, 3, 65, 67], "tolerance": 0.0005, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_m_below_tile_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 63, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 63, 64], "tolerance": 0.0005, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_k_below_tile_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 64, 31],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 64, 31],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 64, 64], "tolerance": 0.0005, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_n_below_tile_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 63, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 64, 63], "tolerance": 0.0005, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_last_A_NaN_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 65, 33],
+          "data": {
+            "kind": "fillFloat32",
+            "sinStep": 0.013,
+            "cosStep": 0.029,
+            "scale": 0.05,
+            "offset": 0.15,
+            "nanStart": 12869.0,
+            "nanCount": 1.0
+          }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 67, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float16",
+          "shape": [2, 3, 65, 67],
+          "tolerance": 0.0005,
+          "relTolerance": 0.001,
+          "allowNaN": true
+        }
+      },
+      "provenance": {
+        "notes": "NaN at the final physical element of a tailed transposed-B broadcast product; checks propagation stays confined to the rows/columns reached by that operand."
+      }
+    },
+    {
+      "name": "broadcast_transb_last_B_NaN_float16",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 67, 33],
+          "data": {
+            "kind": "fillFloat32",
+            "sinStep": 0.019,
+            "cosStep": 0.007,
+            "scale": 0.07,
+            "offset": -0.1,
+            "nanStart": 6632.0,
+            "nanCount": 1.0
+          }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float16",
+          "shape": [2, 3, 65, 67],
+          "tolerance": 0.0005,
+          "relTolerance": 0.001,
+          "allowNaN": true
+        }
+      },
+      "provenance": {
+        "notes": "NaN at the final physical element of a tailed transposed-B broadcast product; checks propagation stays confined to the rows/columns reached by that operand."
+      }
+    },
+    {
+      "name": "broadcast_transb_rank4x3_aligned_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 64, 64], "tolerance": 0.00001, "relTolerance": 0.0001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_rank4x3_tails_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 67, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 65, 67], "tolerance": 0.00001, "relTolerance": 0.0001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_broadcast_a_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 1, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 67, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 65, 67], "tolerance": 0.00001, "relTolerance": 0.0001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_broadcast_b_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [1, 67, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 65, 67], "tolerance": 0.00001, "relTolerance": 0.0001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_rank3x2_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [3, 64, 64], "tolerance": 0.00001, "relTolerance": 0.0001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_rank5x3_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 1, 3, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 67, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": {
+        "Y": { "dtype": "float32", "shape": [2, 1, 3, 65, 67], "tolerance": 0.00001, "relTolerance": 0.0001 }
+      },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_m_below_tile_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 63, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 63, 64], "tolerance": 0.00001, "relTolerance": 0.0001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_k_below_tile_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 64, 31],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 64, 31],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 64, 64], "tolerance": 0.00001, "relTolerance": 0.0001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_n_below_tile_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 63, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 64, 63], "tolerance": 0.00001, "relTolerance": 0.0001 } },
+      "provenance": {
+        "notes": "Transposed-B broadcast addressing: rank, singleton-axis, and tile-boundary coverage; negative alpha and non-unit true transB."
+      }
+    },
+    {
+      "name": "broadcast_transb_last_A_NaN_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 65, 33],
+          "data": {
+            "kind": "fillFloat32",
+            "sinStep": 0.013,
+            "cosStep": 0.029,
+            "scale": 0.05,
+            "offset": 0.15,
+            "nanStart": 12869.0,
+            "nanCount": 1.0
+          }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 67, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.07, "offset": -0.1 }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 3, 65, 67],
+          "tolerance": 0.00001,
+          "relTolerance": 0.0001,
+          "allowNaN": true
+        }
+      },
+      "provenance": {
+        "notes": "NaN at the final physical element of a tailed transposed-B broadcast product; checks propagation stays confined to the rows/columns reached by that operand."
+      }
+    },
+    {
+      "name": "broadcast_transb_last_B_NaN_float32",
+      "attrs": { "transB": 2, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.05, "offset": 0.15 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 67, 33],
+          "data": {
+            "kind": "fillFloat32",
+            "sinStep": 0.019,
+            "cosStep": 0.007,
+            "scale": 0.07,
+            "offset": -0.1,
+            "nanStart": 6632.0,
+            "nanCount": 1.0
+          }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 3, 65, 67],
+          "tolerance": 0.00001,
+          "relTolerance": 0.0001,
+          "allowNaN": true
+        }
+      },
+      "provenance": {
+        "notes": "NaN at the final physical element of a tailed transposed-B broadcast product; checks propagation stays confined to the rows/columns reached by that operand."
+      }
+    },
+    {
+      "name": "sgmat_n_tail_alpha_f32_64x64x65",
+      "attrs": { "alpha": 0.5 },
+      "provenance": {
+        "notes": "The 64-wide subgroup-matrix column tile computes ceilDiv(N, 64) * 64 columns: its trailing tile reads column N - 1 in place of every column at or past N and drops those stores, so a wrong clamp or a missing store guard shows up as a wrong last column or a write past the output. alpha = 0.5 rides the same guarded store, so a tail that skipped the epilogue would show up as an unscaled trailing column."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [64, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [64, 65], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "sgmat_n_tail_transB_f16_64x64x65",
+      "attrs": { "transB": 1 },
+      "provenance": {
+        "notes": "A 65-column float16 transposed-B matrix product checks the final partial column tile, especially the last valid output and guarded stores."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [65, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 65], "tolerance": 0.005, "relTolerance": 0.02 } }
+    },
+    {
+      "name": "sgmat_n_padding_ratio_admits_f32_64x64x32",
+      "provenance": {
+        "notes": "M=64, K=64, N=32 float32 operands check the matrix product at a power-of-two column count; its N=31 sibling is one column narrower."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.017, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [64, 32], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "sgmat_n_padding_ratio_rejects_f32_64x64x31",
+      "provenance": {
+        "notes": "M=64, K=64, N=31 float32 operands check the matrix product one column narrower than its N=32 sibling."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [64, 31],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.017, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [64, 31], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "plain_rank2_tiled_reg_row_tail_m100_k64_n2048",
+      "provenance": {
+        "notes": "M=100 leaves 36 rows past a 64-row boundary (K=64, N=2048), with alpha=0.5 scaling the output; the extra rows must be included correctly in the result."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [100, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.037, "cosStep": 0.061, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [64, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.043, "cosStep": 0.079, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [100, 2048], "tolerance": 0.0001 } }
     }
   ]
 }