Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.FusedMatMul

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
8386dfb
·
verified ·
1 Parent(s): 746b1e6

sync c928d21e6cc1

Browse files
Files changed (8) hide show
  1. README.md +1 -0
  2. build/webgpu/bench.json +66 -0
  3. build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja +14 -14
  4. build/webgpu/manifest.json +98 -4
  5. build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja +43 -20
  6. build/webgpu/metadata.json +8 -7
  7. build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja +34 -0
  8. build/webgpu/test.json +85 -0
README.md CHANGED
@@ -62,6 +62,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
62
  - [`matmul-tiled-general-reg.wgsl.jinja`](build/webgpu/matmul-tiled-general-reg.wgsl.jinja)
63
  - [`matmul-tiled-general.wgsl.jinja`](build/webgpu/matmul-tiled-general.wgsl.jinja)
64
  - [`matmul-vector-matrix-vec4.wgsl.jinja`](build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja)
 
65
 
66
  ## Use with `@huggingface/kernels`
67
 
 
62
  - [`matmul-tiled-general-reg.wgsl.jinja`](build/webgpu/matmul-tiled-general-reg.wgsl.jinja)
63
  - [`matmul-tiled-general.wgsl.jinja`](build/webgpu/matmul-tiled-general.wgsl.jinja)
64
  - [`matmul-vector-matrix-vec4.wgsl.jinja`](build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja)
65
+ - [`reduce-axis0-splitk-combine.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja)
66
 
67
  ## Use with `@huggingface/kernels`
68
 
build/webgpu/bench.json CHANGED
@@ -212,6 +212,72 @@
212
  },
213
  "outputs": { "Y": { "shape": [1, 8, 512, 512], "dtype": "float16", "dist": "empty" } },
214
  "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 512 * 512 * 2048" }] }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
215
  }
216
  ]
217
  }
 
212
  },
213
  "outputs": { "Y": { "shape": [1, 8, 512, 512], "dtype": "float16", "dist": "empty" } },
214
  "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 512 * 512 * 2048" }] }
215
+ },
216
+ {
217
+ "name": "fusedmatmul-f16-band-m8-k2560-n4096",
218
+ "preset": "model",
219
+ "attrs": { "alpha": 0.5 },
220
+ "inputs": {
221
+ "A": { "shape": [8, 2560], "dtype": "float16", "dist": "normal", "seed": 750, "scale": 0.1 },
222
+ "B": { "shape": [2560, 4096], "dtype": "float16", "dist": "normal", "seed": 751, "scale": 0.1 }
223
+ },
224
+ "outputs": { "Y": { "shape": [8, 4096], "dtype": "float16", "dist": "empty" } },
225
+ "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 2560 * 4096" }] }
226
+ },
227
+ {
228
+ "name": "fusedmatmul-f16-band-m16-k2560-n4096",
229
+ "preset": "model",
230
+ "attrs": { "alpha": 0.5 },
231
+ "inputs": {
232
+ "A": { "shape": [16, 2560], "dtype": "float16", "dist": "normal", "seed": 752, "scale": 0.1 },
233
+ "B": { "shape": [2560, 4096], "dtype": "float16", "dist": "normal", "seed": 753, "scale": 0.1 }
234
+ },
235
+ "outputs": { "Y": { "shape": [16, 4096], "dtype": "float16", "dist": "empty" } },
236
+ "bench": { "metrics": [{ "type": "gflops", "value": "2 * 16 * 2560 * 4096" }] }
237
+ },
238
+ {
239
+ "name": "fusedmatmul-f16-band-m32-k2560-n4096",
240
+ "preset": "model",
241
+ "attrs": { "alpha": 0.5 },
242
+ "inputs": {
243
+ "A": { "shape": [32, 2560], "dtype": "float16", "dist": "normal", "seed": 754, "scale": 0.1 },
244
+ "B": { "shape": [2560, 4096], "dtype": "float16", "dist": "normal", "seed": 755, "scale": 0.1 }
245
+ },
246
+ "outputs": { "Y": { "shape": [32, 4096], "dtype": "float16", "dist": "empty" } },
247
+ "bench": { "metrics": [{ "type": "gflops", "value": "2 * 32 * 2560 * 4096" }] }
248
+ },
249
+ {
250
+ "name": "fusedmatmul-f32-band-m16-k2560-n4096",
251
+ "preset": "model",
252
+ "attrs": { "alpha": 0.5 },
253
+ "inputs": {
254
+ "A": { "shape": [16, 2560], "dtype": "float32", "dist": "normal", "seed": 756, "scale": 0.1 },
255
+ "B": { "shape": [2560, 4096], "dtype": "float32", "dist": "normal", "seed": 757, "scale": 0.1 }
256
+ },
257
+ "outputs": { "Y": { "shape": [16, 4096], "dtype": "float32", "dist": "empty" } },
258
+ "bench": { "metrics": [{ "type": "gflops", "value": "2 * 16 * 2560 * 4096" }] }
259
+ },
260
+ {
261
+ "name": "fusedmatmul-f16-narrow-m16-k8192-n256",
262
+ "preset": "model",
263
+ "attrs": { "alpha": 1 },
264
+ "inputs": {
265
+ "A": { "shape": [16, 8192], "dtype": "float16", "dist": "normal", "seed": 758, "scale": 0.1 },
266
+ "B": { "shape": [8192, 256], "dtype": "float16", "dist": "normal", "seed": 759, "scale": 0.1 }
267
+ },
268
+ "outputs": { "Y": { "shape": [16, 256], "dtype": "float16", "dist": "empty" } },
269
+ "bench": { "metrics": [{ "type": "gflops", "value": "2 * 16 * 8192 * 256" }] }
270
+ },
271
+ {
272
+ "name": "fusedmatmul-f16-specdec-scores-transB-8xm4-k64-n512",
273
+ "preset": "model",
274
+ "attrs": { "alpha": 0.125, "transB": 1 },
275
+ "inputs": {
276
+ "A": { "shape": [8, 4, 64], "dtype": "float16", "dist": "normal", "seed": 760, "scale": 0.1 },
277
+ "B": { "shape": [8, 512, 64], "dtype": "float16", "dist": "normal", "seed": 761, "scale": 0.1 }
278
+ },
279
+ "outputs": { "Y": { "shape": [8, 4, 512], "dtype": "float16", "dist": "empty" } },
280
+ "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 4 * 64 * 512" }] }
281
  }
282
  ]
283
  }
build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja CHANGED
@@ -128,14 +128,14 @@ fn main(
128
 
129
  for (var step = 0u; step < TILE_K; step = step + 8u) {
130
  let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
131
- var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset, false, TILE_K);
132
- var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, TILE_K);
133
 
134
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
135
- var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset, true, TILE_K);
136
- var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
137
- var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
138
- var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
139
 
140
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
141
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -154,10 +154,10 @@ fn main(
154
  // give the reads visibility of the store and stop the second row-group's store from
155
  // clobbering the first's still-in-flight readback when a partial final M-tile
156
  // diverges storeOutput's guard. Without both barriers the last valid row can be corrupted.
157
- subgroupMatrixStore(&scratch[subtile_id][0], 0u, matC00, false, 8u);
158
- subgroupMatrixStore(&scratch[subtile_id][1], 0u, matC01, false, 8u);
159
- subgroupMatrixStore(&scratch[subtile_id][2], 0u, matC02, false, 8u);
160
- subgroupMatrixStore(&scratch[subtile_id][3], 0u, matC03, false, 8u);
161
  workgroupBarrier();
162
  let row = sg_id / 4u;
163
  let col = (sg_id % 4u) * 2u;
@@ -166,10 +166,10 @@ fn main(
166
  storeOutput(matrix_c_offset, row, col, subtile_id, row_limit);
167
  workgroupBarrier();
168
 
169
- subgroupMatrixStore(&scratch[subtile_id][0], 0u, matC10, false, 8u);
170
- subgroupMatrixStore(&scratch[subtile_id][1], 0u, matC11, false, 8u);
171
- subgroupMatrixStore(&scratch[subtile_id][2], 0u, matC12, false, 8u);
172
- subgroupMatrixStore(&scratch[subtile_id][3], 0u, matC13, false, 8u);
173
  workgroupBarrier();
174
  matrix_c_offset = matrix_c_offset + 8u * N;
175
  row_limit = i32(M) - i32(a_global_base + base_A + 8u);
 
128
 
129
  for (var step = 0u; step < TILE_K; step = step + 8u) {
130
  let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
131
+ var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
132
+ var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
133
 
134
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
135
+ var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
136
+ var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
137
+ var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
138
+ var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
139
 
140
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
141
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
 
154
  // give the reads visibility of the store and stop the second row-group's store from
155
  // clobbering the first's still-in-flight readback when a partial final M-tile
156
  // diverges storeOutput's guard. Without both barriers the last valid row can be corrupted.
157
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC00, 8u);
158
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][1], 0u, matC01, 8u);
159
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][2], 0u, matC02, 8u);
160
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][3], 0u, matC03, 8u);
161
  workgroupBarrier();
162
  let row = sg_id / 4u;
163
  let col = (sg_id % 4u) * 2u;
 
166
  storeOutput(matrix_c_offset, row, col, subtile_id, row_limit);
167
  workgroupBarrier();
168
 
169
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC10, 8u);
170
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][1], 0u, matC11, 8u);
171
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][2], 0u, matC12, 8u);
172
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][3], 0u, matC13, 8u);
173
  workgroupBarrier();
174
  matrix_c_offset = matrix_c_offset + 8u * N;
175
  row_limit = i32(M) - i32(a_global_base + base_A + 8u);
build/webgpu/manifest.json CHANGED
@@ -30,13 +30,35 @@
30
  "B": { "kind": "tensor", "semantic": "B", "role": "input" },
31
  "Y": { "kind": "tensor", "semantic": "Y", "role": "output" }
32
  },
33
- "tunables": { "TILED_REG_MIN_WORKGROUPS": 64, "GEMV_TARGET_BLOCKS": 512 },
 
 
 
 
 
 
 
34
  "bindingSets": {
35
  "matrix": [
36
  { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
37
  { "name": "b", "arg": "B", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
38
  { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" }
39
  ],
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40
  "gemvF32": [
41
  { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
42
  {
@@ -73,7 +95,19 @@
73
  "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
74
  "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",
75
  "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
76
- "wave32Effective": "wave32Adapter or pinSubgroupSize32"
 
 
 
 
 
 
 
 
 
 
 
 
77
  },
78
  "variants": [
79
  {
@@ -91,6 +125,66 @@
91
  }
92
  ]
93
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
94
  {
95
  "id": "subgroup_matrix_tail_broadcast",
96
  "priority": 11,
@@ -98,7 +192,7 @@
98
  "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
99
  "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
100
  },
101
- "when": ["dtypes.T == \"f16\"", "f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "((ranks.A == 2 and ranks.B == 2 and ranks.Y == 2) or (ranks.A == 4 and ranks.B == 3 and ranks.Y == 4 and dim(shapes.Y, 0) == dim(shapes.A, 0) and (dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1) and dim(shapes.Y, 1) == max(dim(shapes.A, 1), dim(shapes.B, 0))))", "dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 2)", "dim(shapes.A, ranks.A - 2) >= 32", "dim(shapes.A, ranks.A - 1) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "dim(shapes.Y, ranks.Y - 2) == dim(shapes.A, ranks.A - 2)", "dim(shapes.Y, ranks.Y - 1) == dim(shapes.B, ranks.B - 1)", "ceil(dim(shapes.B, ranks.B - 1) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, ranks.A - 2) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"],
102
  "constants": {
103
  "hasBias": false,
104
  "usesF16": true,
@@ -137,7 +231,7 @@
137
  { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
138
  ]
139
  },
140
- "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))) >= 32", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) % 32 == 0", "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) % 64 == 0", "(ranks.A == 2 or (ranks.A == 3 and (dim(shapes.A, 1) if attrs.transBatchA != 0 else dim(shapes.A, 0)) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.B, 0)) or (ranks.A == 4 and dim(shapes.A, 0) == dim(shapes.B, 0) and dim(shapes.A, 1) == dim(shapes.B, 1) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.A, 1)))", "dim(shapes.Y, ranks.Y - 2) == (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))", "dim(shapes.Y, ranks.Y - 1) == (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))", "ceil((dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) * (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"],
141
  "constants": {
142
  "usesF16": "dtypes.T == \"f16\"",
143
  "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
 
30
  "B": { "kind": "tensor", "semantic": "B", "role": "input" },
31
  "Y": { "kind": "tensor", "semantic": "Y", "role": "output" }
32
  },
33
+ "tunables": {
34
+ "TILED_REG_MIN_WORKGROUPS": 64,
35
+ "GEMV_TARGET_BLOCKS": 512,
36
+ "SUBGROUP_MATRIX_MIN_M": 2,
37
+ "SUBGROUP_MATRIX_SPLITK_TARGET_WGS": 512,
38
+ "SUBGROUP_MATRIX_SPLITK_MIN_K": 1024,
39
+ "SUBGROUP_MATRIX_SPLITK_MAX_TILES": 128
40
+ },
41
  "bindingSets": {
42
  "matrix": [
43
  { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
44
  { "name": "b", "arg": "B", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
45
  { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" }
46
  ],
47
+ "sgmatSplitPartial": [
48
+ { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
49
+ { "name": "b", "arg": "B", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
50
+ { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" }
51
+ ],
52
+ "sgmatSplitCombine": [
53
+ { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
54
+ { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
55
+ {
56
+ "name": "params",
57
+ "semantic": "kernel.params",
58
+ "buffer": { "type": "uniform" },
59
+ "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "numel(shapes.Y)" }] }
60
+ }
61
+ ],
62
  "gemvF32": [
63
  { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
64
  {
 
95
  "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
96
  "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",
97
  "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
98
+ "wave32Effective": "wave32Adapter or pinSubgroupSize32",
99
+ "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
100
+ "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)",
101
+ "subgroupMatrixResourcesFit": "128 <= deviceWorkgroupCap and ((32 * 32 + 64 * 32) * dtypeBytes(dtypes.T) + 4 * 4 * 64 * 4) <= device.limits.maxComputeWorkgroupStorageSize",
102
+ "sgmatSplitKDepth": "dim(shapes.A, ranks.A - 1)",
103
+ "sgmatOutTiles": "ceilDiv(dim(shapes.A, 0), 32) * ceilDiv(dim(shapes.B, 1), 64) if fusedSgmatRank2Ok else 1",
104
+ "sgmatSplitKWant": "ceilDiv(tunables.SUBGROUP_MATRIX_SPLITK_TARGET_WGS, sgmatOutTiles)",
105
+ "sgmatSplitK32Ok": "sgmatSplitKDepth % 1024 == 0",
106
+ "sgmatSplitK16Ok": "sgmatSplitKDepth % 512 == 0",
107
+ "sgmatSplitK8Ok": "sgmatSplitKDepth % 256 == 0",
108
+ "sgmatSplitK4Ok": "sgmatSplitKDepth % 128 == 0",
109
+ "sgmatSplitK2Ok": "sgmatSplitKDepth % 64 == 0",
110
+ "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))))"
111
  },
112
  "variants": [
113
  {
 
125
  }
126
  ]
127
  },
128
+ {
129
+ "id": "subgroup_matrix_splitk",
130
+ "description": "Partitions the K reduction across workgroups while retaining subgroup-matrix tiles, then combines float32 partials that already carry alpha. It serves small-M rank-two products whose unsplit tile grids cannot occupy the device.",
131
+ "priority": 12,
132
+ "requires": {
133
+ "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
134
+ "subgroupMatrixConfigs": [
135
+ { "componentType": "f16", "M": 8, "N": 8, "K": 8 },
136
+ { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
137
+ ]
138
+ },
139
+ "when": ["(dtypes.T == \"f16\" or dtypes.T == \"f32\") and f16Ok(dtypes.T)", "fusedSgmatRank2Ok", "dim(shapes.A, 0) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, 1) >= tunables.SUBGROUP_MATRIX_SPLITK_MIN_K", "dim(shapes.B, 1) % 64 == 0", "sgmatSplitK >= 2", "sgmatOutTiles < tunables.SUBGROUP_MATRIX_SPLITK_MAX_TILES", "sgmatSplitK * numel(shapes.Y) * 4 <= device.limits.maxStorageBufferBindingSize", "sgmatSplitK <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(dim(shapes.Y, 1), 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(dim(shapes.Y, 0), 32) <= device.limits.maxComputeWorkgroupsPerDimension", "subgroupMatrixResourcesFit", "wave32Effective"],
140
+ "constants": {
141
+ "usesF16": "dtypes.T == \"f16\"",
142
+ "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
143
+ "scalar": "dtypes.T",
144
+ "hasBias": false,
145
+ "generalAddressing": true,
146
+ "tailSafe": false,
147
+ "outputBuffer": "\"partials\"",
148
+ "outScalar": "\"f32\"",
149
+ "alpha": "attrs.alpha",
150
+ "M": "dim(shapes.A, 0)",
151
+ "K": "dim(shapes.A, 1)",
152
+ "N": "dim(shapes.B, 1)",
153
+ "batchCount": 1,
154
+ "splitK": "sgmatSplitK",
155
+ "kPerSplit": "dim(shapes.A, 1) / sgmatSplitK",
156
+ "split": "sgmatSplitK",
157
+ "workgroupSize": 256
158
+ },
159
+ "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[sgmatSplitK * numel(shapes.Y)]" }],
160
+ "passes": [
161
+ {
162
+ "id": "partial",
163
+ "name": "FusedMatMul.SubgroupMatrixSplitK",
164
+ "source": {
165
+ "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
166
+ "inputs": {
167
+ "aShape": ["dim(shapes.A, 0)", "dim(shapes.A, 1)"],
168
+ "bShape": ["dim(shapes.B, 0)", "dim(shapes.B, 1)"],
169
+ "aRank": 2,
170
+ "bRank": 2
171
+ }
172
+ },
173
+ "bindings": "sgmatSplitPartial",
174
+ "dispatch": { "x": "ceilDiv(dim(shapes.Y, 1), 64)", "y": "ceilDiv(dim(shapes.Y, 0), 32)", "z": "sgmatSplitK" }
175
+ },
176
+ {
177
+ "id": "combine",
178
+ "name": "FusedMatMul.SubgroupMatrixSplitKCombine",
179
+ "source": {
180
+ "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
181
+ "inputs": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"", "intMode": false }
182
+ },
183
+ "bindings": "sgmatSplitCombine",
184
+ "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": 256 }
185
+ }
186
+ ]
187
+ },
188
  {
189
  "id": "subgroup_matrix_tail_broadcast",
190
  "priority": 11,
 
192
  "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
193
  "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
194
  },
195
+ "when": ["dtypes.T == \"f16\"", "f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "((ranks.A == 2 and ranks.B == 2 and ranks.Y == 2) or (ranks.A == 4 and ranks.B == 3 and ranks.Y == 4 and dim(shapes.Y, 0) == dim(shapes.A, 0) and (dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1) and dim(shapes.Y, 1) == max(dim(shapes.A, 1), dim(shapes.B, 0))))", "dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 2)", "dim(shapes.A, ranks.A - 2) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, ranks.A - 1) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "dim(shapes.Y, ranks.Y - 2) == dim(shapes.A, ranks.A - 2)", "dim(shapes.Y, ranks.Y - 1) == dim(shapes.B, ranks.B - 1)", "ceil(dim(shapes.B, ranks.B - 1) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, ranks.A - 2) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"],
196
  "constants": {
197
  "hasBias": false,
198
  "usesF16": true,
 
231
  { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
232
  ]
233
  },
234
+ "when": ["f16Ok(dtypes.T)", "attrs.transBatchA == 0 or (attrs.transA == 0 and ranks.A == 3)", "attrs.transBatchB == 0", "ranks.A >= 2", "ranks.B == ranks.A", "ranks.Y == ranks.A", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) == (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else dim(shapes.B, ranks.B - 2))", "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) >= tunables.SUBGROUP_MATRIX_MIN_M", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) % 32 == 0", "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) % 64 == 0", "(ranks.A == 2 or (ranks.A == 3 and (dim(shapes.A, 1) if attrs.transBatchA != 0 else dim(shapes.A, 0)) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.B, 0)) or (ranks.A == 4 and dim(shapes.A, 0) == dim(shapes.B, 0) and dim(shapes.A, 1) == dim(shapes.B, 1) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.A, 1)))", "dim(shapes.Y, ranks.Y - 2) == (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))", "dim(shapes.Y, ranks.Y - 1) == (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))", "ceil((dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) * (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"],
235
  "constants": {
236
  "usesF16": "dtypes.T == \"f16\"",
237
  "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja CHANGED
@@ -17,6 +17,7 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
17
  {% set accScalar = "f32" %}
18
  {% set GENERAL = generalAddressing is defined and generalAddressing %}
19
  {% set TAIL = tailSafe is defined and tailSafe %}
 
20
  {% set OUT = outputBuffer if outputBuffer is defined else "c" %}
21
  {% set OUT_SCALAR = outScalar if outScalar is defined else T %}
22
  {% if GENERAL %}
@@ -35,6 +36,10 @@ const M: u32 = {{ M }}u;
35
  const K: u32 = {{ K }}u;
36
  const N: u32 = {{ N }}u;
37
  const BATCH_COUNT: u32 = {{ batchCount if batchCount is defined else 1 }}u;
 
 
 
 
38
  {% if GENERAL %}
39
  const A_M_STRIDE: u32 = {{ aMStride }}u;
40
  const A_K_STRIDE: u32 = {{ aKStride }}u;
@@ -216,8 +221,16 @@ fn main(
216
  // exactly once at workgroup_id.y/workgroup_id.z.
217
  let row_tile_stride = num_wg.y * TILE_ROWS;
218
  for (var a_global_base = workgroup_id.y * TILE_ROWS; a_global_base < M; a_global_base += row_tile_stride) {
 
 
 
 
 
 
 
219
  // workgroup_size.z = 1, so num_wg.z is the dispatch stride over the batch axis.
220
  for (var batch = workgroup_id.z; batch < BATCH_COUNT; batch += num_wg.z) {
 
221
  {% if GENERAL %}
222
  {% set hasBatchCoord = namespace(value=false) %}
223
  {% for i in range(batchRank) %}
@@ -254,7 +267,11 @@ fn main(
254
  {% endif %}
255
  {% endif %}
256
  {% endfor %}
 
 
 
257
  let c_base = batch * C_BATCH_STRIDE;
 
258
  {% else %}
259
  let a_base = batch * A_BATCH_STRIDE;
260
  let b_base = batch * B_BATCH_STRIDE;
@@ -270,21 +287,27 @@ fn main(
270
  var matC12: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
271
  var matC13: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
272
 
 
 
 
 
 
273
  for (var kidx = 0u; kidx < {% if GENERAL and TAIL %}K_FULL{% else %}K{% endif %}; kidx = kidx + TILE_K) {
 
274
  loadSHMA(a_base, a_global_base, kidx, local_idx / 4u, local_idx % 4u);
275
  loadSHMB(b_base, b_global_base, kidx, local_idx / 2u, local_idx % 2u);
276
  workgroupBarrier();
277
 
278
  for (var step = 0u; step < TILE_K; step = step + 8u) {
279
  let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
280
- var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset, false, TILE_K);
281
- var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, TILE_K);
282
 
283
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
284
- var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset, true, TILE_K);
285
- var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
286
- var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
287
- var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
288
 
289
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
290
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -305,14 +328,14 @@ fn main(
305
 
306
  for (var step = 0u; step < TILE_K; step = step + 8u) {
307
  let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
308
- var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset, false, TILE_K);
309
- var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, TILE_K);
310
 
311
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
312
- var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset, true, TILE_K);
313
- var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
314
- var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
315
- var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
316
 
317
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
318
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -332,10 +355,10 @@ fn main(
332
  // give the reads visibility of the store AND stop the second row-group's store from
333
  // clobbering the first's still-in-flight readback when a partial final M-tile
334
  // diverges storeOutput's guard. Without both barriers the last valid row can be corrupted.
335
- subgroupMatrixStore(&scratch[subtile_id][0], 0u, matC00, false, 8u);
336
- subgroupMatrixStore(&scratch[subtile_id][1], 0u, matC01, false, 8u);
337
- subgroupMatrixStore(&scratch[subtile_id][2], 0u, matC02, false, 8u);
338
- subgroupMatrixStore(&scratch[subtile_id][3], 0u, matC03, false, 8u);
339
  workgroupBarrier();
340
  let row = sg_id / 4u;
341
  let col = (sg_id % 4u) * 2u;
@@ -345,10 +368,10 @@ fn main(
345
  storeOutput(matrix_c_offset{% if needsColBase %}, col_base{% endif %}, row, col, subtile_id, row_limit);
346
  workgroupBarrier();
347
 
348
- subgroupMatrixStore(&scratch[subtile_id][0], 0u, matC10, false, 8u);
349
- subgroupMatrixStore(&scratch[subtile_id][1], 0u, matC11, false, 8u);
350
- subgroupMatrixStore(&scratch[subtile_id][2], 0u, matC12, false, 8u);
351
- subgroupMatrixStore(&scratch[subtile_id][3], 0u, matC13, false, 8u);
352
  workgroupBarrier();
353
  matrix_c_offset = matrix_c_offset + 8u * N;
354
  row_limit = i32(M) - i32(a_global_base + base_A + 8u);
 
17
  {% set accScalar = "f32" %}
18
  {% set GENERAL = generalAddressing is defined and generalAddressing %}
19
  {% set TAIL = tailSafe is defined and tailSafe %}
20
+ {% set SPLIT_K = splitK if splitK is defined else 1 %}
21
  {% set OUT = outputBuffer if outputBuffer is defined else "c" %}
22
  {% set OUT_SCALAR = outScalar if outScalar is defined else T %}
23
  {% if GENERAL %}
 
36
  const K: u32 = {{ K }}u;
37
  const N: u32 = {{ N }}u;
38
  const BATCH_COUNT: u32 = {{ batchCount if batchCount is defined else 1 }}u;
39
+ {% if SPLIT_K > 1 %}
40
+ const SPLIT_K: u32 = {{ SPLIT_K }}u;
41
+ const K_PER_SPLIT: u32 = {{ kPerSplit }}u;
42
+ {% endif %}
43
  {% if GENERAL %}
44
  const A_M_STRIDE: u32 = {{ aMStride }}u;
45
  const A_K_STRIDE: u32 = {{ aKStride }}u;
 
221
  // exactly once at workgroup_id.y/workgroup_id.z.
222
  let row_tile_stride = num_wg.y * TILE_ROWS;
223
  for (var a_global_base = workgroup_id.y * TILE_ROWS; a_global_base < M; a_global_base += row_tile_stride) {
224
+ {% if SPLIT_K > 1 %}
225
+ // Split-K maps z to (batch, K segment), multiplying occupancy for narrow
226
+ // matrices while every segment still owns a disjoint contiguous K range.
227
+ for (var batch_split = workgroup_id.z; batch_split < BATCH_COUNT * SPLIT_K; batch_split += num_wg.z) {
228
+ let batch = batch_split / SPLIT_K;
229
+ let split_id = batch_split - batch * SPLIT_K;
230
+ {% else %}
231
  // workgroup_size.z = 1, so num_wg.z is the dispatch stride over the batch axis.
232
  for (var batch = workgroup_id.z; batch < BATCH_COUNT; batch += num_wg.z) {
233
+ {% endif %}
234
  {% if GENERAL %}
235
  {% set hasBatchCoord = namespace(value=false) %}
236
  {% for i in range(batchRank) %}
 
267
  {% endif %}
268
  {% endif %}
269
  {% endfor %}
270
+ {% if SPLIT_K > 1 %}
271
+ let c_base = (batch * SPLIT_K + split_id) * C_BATCH_STRIDE;
272
+ {% else %}
273
  let c_base = batch * C_BATCH_STRIDE;
274
+ {% endif %}
275
  {% else %}
276
  let a_base = batch * A_BATCH_STRIDE;
277
  let b_base = batch * B_BATCH_STRIDE;
 
287
  var matC12: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
288
  var matC13: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
289
 
290
+ {% if SPLIT_K > 1 %}
291
+ let k_begin = split_id * K_PER_SPLIT;
292
+ let k_end = min(k_begin + K_PER_SPLIT, K);
293
+ for (var kidx = k_begin; kidx < k_end; kidx = kidx + TILE_K) {
294
+ {% else %}
295
  for (var kidx = 0u; kidx < {% if GENERAL and TAIL %}K_FULL{% else %}K{% endif %}; kidx = kidx + TILE_K) {
296
+ {% endif %}
297
  loadSHMA(a_base, a_global_base, kidx, local_idx / 4u, local_idx % 4u);
298
  loadSHMB(b_base, b_global_base, kidx, local_idx / 2u, local_idx % 2u);
299
  workgroupBarrier();
300
 
301
  for (var step = 0u; step < TILE_K; step = step + 8u) {
302
  let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
303
+ var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
304
+ var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
305
 
306
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
307
+ var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
308
+ var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
309
+ var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
310
+ var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
311
 
312
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
313
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
 
328
 
329
  for (var step = 0u; step < TILE_K; step = step + 8u) {
330
  let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
331
+ var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
332
+ var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
333
 
334
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
335
+ var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
336
+ var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
337
+ var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
338
+ var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
339
 
340
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
341
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
 
355
  // give the reads visibility of the store AND stop the second row-group's store from
356
  // clobbering the first's still-in-flight readback when a partial final M-tile
357
  // diverges storeOutput's guard. Without both barriers the last valid row can be corrupted.
358
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC00, 8u);
359
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][1], 0u, matC01, 8u);
360
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][2], 0u, matC02, 8u);
361
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][3], 0u, matC03, 8u);
362
  workgroupBarrier();
363
  let row = sg_id / 4u;
364
  let col = (sg_id % 4u) * 2u;
 
368
  storeOutput(matrix_c_offset{% if needsColBase %}, col_base{% endif %}, row, col, subtile_id, row_limit);
369
  workgroupBarrier();
370
 
371
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC10, 8u);
372
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][1], 0u, matC11, 8u);
373
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][2], 0u, matC12, 8u);
374
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][3], 0u, matC13, 8u);
375
  workgroupBarrier();
376
  matrix_c_offset = matrix_c_offset + 8u * N;
377
  row_limit = i32(M) - i32(a_global_base + base_A + 8u);
build/webgpu/metadata.json CHANGED
@@ -1,22 +1,23 @@
1
  {
2
  "name": "com.microsoft.FusedMatMul",
3
- "id": "_com_microsoft_fusedmatmul_webgpu_80c34de",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "uw58QO9KL2kJYYMWMZJSUGIwBmsy6bLImlj0E1NNCyw=",
11
- "fused-matmul-subgroup-matrix.wgsl.jinja": "d8LaODNB5B+qCRM5SSDSbasoJTBoTkjR8BtwuO8AW6U=",
12
- "manifest.json": "F7nDLMjrVNWkpIGc9rUB5ZrhRoW06xk+PPX+M6X7XCA=",
13
- "matmul-subgroup-matrix-ext.wgsl.jinja": "Q/5RQ4YPRDz/EiGmw1HigsSCohWzsx6muR4ghzigfWs=",
14
  "matmul-tiled-general-reg.wgsl.jinja": "8E6pn6tg6t7LuoU6ktkjL5AZOgrd0HHkFRe8VJC7jwk=",
15
  "matmul-tiled-general.wgsl.jinja": "pqYgOQCvHqeTthCEShE8dbN1GJjpFz3LD1K3sBYARMc=",
16
  "matmul-vector-matrix-vec4.wgsl.jinja": "9syYZdAgoDnCn5rurCHP9tHtlwm1IXwKvPoj48hH3+E=",
17
- "test.json": "Au5pkdVg7fcbm2EDBuhQvc4CPiri5lrQs394cLC5iL0="
 
18
  }
19
  },
20
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
21
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.FusedMatMul" }
22
  }
 
1
  {
2
  "name": "com.microsoft.FusedMatMul",
3
+ "id": "_com_microsoft_fusedmatmul_webgpu_a7be42e",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "FKzhB1qpaKYuQ28OUS2fiiJaDroyLDXWubS0iFa3Fyw=",
11
+ "fused-matmul-subgroup-matrix.wgsl.jinja": "tZdeSq2o0dor9WP7R5Tf5Yzxm8DXHapNXdGci8SH4wI=",
12
+ "manifest.json": "TpmC/4RDrMS30inZSCKjBEw3wo+UAldsvOfUrpKuZn0=",
13
+ "matmul-subgroup-matrix-ext.wgsl.jinja": "W5PywIfrssv2VQhuCG1lu3LjC3R7VpxGRnkHwMX050Q=",
14
  "matmul-tiled-general-reg.wgsl.jinja": "8E6pn6tg6t7LuoU6ktkjL5AZOgrd0HHkFRe8VJC7jwk=",
15
  "matmul-tiled-general.wgsl.jinja": "pqYgOQCvHqeTthCEShE8dbN1GJjpFz3LD1K3sBYARMc=",
16
  "matmul-vector-matrix-vec4.wgsl.jinja": "9syYZdAgoDnCn5rurCHP9tHtlwm1IXwKvPoj48hH3+E=",
17
+ "reduce-axis0-splitk-combine.wgsl.jinja": "uvQNNkj1ezvpKVhYyLUAAfGVLOk/i7VnAiiNSYYz2uk=",
18
+ "test.json": "Nv5BBqSF2s59CpK7a6Zmu1cc32XOn+PNz8J/MiLO0uo="
19
  }
20
  },
21
+ "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
22
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.FusedMatMul" }
23
  }
build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Pass 2 of the split-K column-wise reduction. One invocation per output column
2
+ // folds the segment partials and applies the selected reduction's final step.
3
+ // Segments are folded in ascending order for deterministic results. This order
4
+ // differs from the single-pass reduction but remains within the f32 tolerance.
5
+ {% if source.outputF16 %}
6
+ enable f16;
7
+ {% endif %}
8
+ {{ env.wgsl.resourceDeclarations }}
9
+ /* Exact max/min reduction identities. WGSL rejects infinity during constant
10
+ * evaluation, so f32 identities are constructed at runtime from their
11
+ * IEEE-754 bit patterns. */
12
+
13
+
14
+ const WG: u32 = {{ workgroupSize }}u;
15
+ const SPLIT: u32 = {{ split }}u;
16
+
17
+ @compute @workgroup_size(WG, 1, 1)
18
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
19
+ @builtin(num_workgroups) nwg: vec3<u32>) {
20
+ let stride = nwg.x * WG;
21
+ let start = (gid.y * nwg.x * WG) + gid.x;
22
+ for (var col = start; col < params.cols; col = col + stride) {
23
+ var total = 0.0;
24
+ for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
25
+ let p = partials[seg * params.cols + col];
26
+ total = total + p;
27
+ }
28
+ {% if source.outputF16 %}
29
+ y[col] = f16(total);
30
+ {% else %}
31
+ y[col] = total;
32
+ {% endif %}
33
+ }
34
+ }
build/webgpu/test.json CHANGED
@@ -1415,6 +1415,91 @@
1415
  }
1416
  }
1417
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1418
  }
1419
  ]
1420
  }
 
1415
  }
1416
  }
1417
  }
1418
+ },
1419
+ {
1420
+ "name": "subgroup_matrix_band_m8_f16",
1421
+ "attrs": { "alpha": 2 },
1422
+ "inputs": {
1423
+ "A": {
1424
+ "dtype": "float16",
1425
+ "shape": [8, 64],
1426
+ "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
1427
+ },
1428
+ "B": {
1429
+ "dtype": "float16",
1430
+ "shape": [64, 64],
1431
+ "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029 }
1432
+ }
1433
+ },
1434
+ "outputs": { "Y": { "dtype": "float16", "shape": [8, 64], "tolerance": 0.005 } }
1435
+ },
1436
+ {
1437
+ "name": "subgroup_matrix_splitk_m_tail_f16",
1438
+ "attrs": { "alpha": 2 },
1439
+ "inputs": {
1440
+ "A": {
1441
+ "dtype": "float16",
1442
+ "shape": [10, 2048],
1443
+ "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
1444
+ },
1445
+ "B": {
1446
+ "dtype": "float16",
1447
+ "shape": [2048, 256],
1448
+ "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.031 }
1449
+ }
1450
+ },
1451
+ "outputs": { "Y": { "dtype": "float16", "shape": [10, 256], "tolerance": 0.005 } }
1452
+ },
1453
+ {
1454
+ "name": "subgroup_matrix_splitk_alpha_scaled_f32",
1455
+ "attrs": { "alpha": 1.5 },
1456
+ "inputs": {
1457
+ "A": {
1458
+ "dtype": "float32",
1459
+ "shape": [16, 1024],
1460
+ "data": { "kind": "fillFloat32", "sinStep": 0.03, "cosStep": 0.07, "scale": 0.2 }
1461
+ },
1462
+ "B": {
1463
+ "dtype": "float32",
1464
+ "shape": [1024, 128],
1465
+ "data": { "kind": "fillFloat32", "sinStep": 0.041, "cosStep": 0.089, "scale": 0.2 }
1466
+ }
1467
+ },
1468
+ "outputs": { "Y": { "dtype": "float32", "shape": [16, 128], "tolerance": 0.0002 } }
1469
+ },
1470
+ {
1471
+ "name": "subgroup_matrix_batched_transB_small_m_f16",
1472
+ "attrs": { "transB": 1, "alpha": 0.25 },
1473
+ "inputs": {
1474
+ "A": {
1475
+ "dtype": "float16",
1476
+ "shape": [2, 4, 32],
1477
+ "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2 }
1478
+ },
1479
+ "B": {
1480
+ "dtype": "float16",
1481
+ "shape": [2, 64, 32],
1482
+ "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
1483
+ }
1484
+ },
1485
+ "outputs": { "Y": { "dtype": "float16", "shape": [2, 4, 64], "tolerance": 0.01 } }
1486
+ },
1487
+ {
1488
+ "name": "subgroup_matrix_transA_small_m_f16",
1489
+ "attrs": { "transA": 1, "alpha": 3 },
1490
+ "inputs": {
1491
+ "A": {
1492
+ "dtype": "float16",
1493
+ "shape": [64, 8],
1494
+ "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
1495
+ },
1496
+ "B": {
1497
+ "dtype": "float16",
1498
+ "shape": [64, 64],
1499
+ "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029 }
1500
+ }
1501
+ },
1502
+ "outputs": { "Y": { "dtype": "float16", "shape": [8, 64], "tolerance": 0.005 } }
1503
  }
1504
  ]
1505
  }