Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.MatMulBnb4

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
902be69
·
verified ·
1 Parent(s): adee57e

sync 6fdf6301e2bb

Browse files
Files changed (8) hide show
  1. README.md +2 -2
  2. build/webgpu/bench.json +1 -1
  3. build/webgpu/manifest.json +24 -104
  4. build/webgpu/matmul-bnb4-gemv.wgsl.jinja +2 -4
  5. build/webgpu/matmul-bnb4-sgmat.wgsl.jinja +2 -7
  6. build/webgpu/matmul-bnb4-tiled.wgsl.jinja +2 -4
  7. build/webgpu/matmul-bnb4.wgsl.jinja +8 -8
  8. build/webgpu/metadata.json +9 -9
README.md CHANGED
@@ -68,7 +68,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
68
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
69
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
70
  - [`test.json`](build/webgpu/test.json) — correctness cases
71
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
72
  - [`cast-scalar-x4.wgsl.jinja`](build/webgpu/cast-scalar-x4.wgsl.jinja)
73
  - [`matmul-bnb4-gemv.wgsl.jinja`](build/webgpu/matmul-bnb4-gemv.wgsl.jinja)
74
  - [`matmul-bnb4-sgmat.wgsl.jinja`](build/webgpu/matmul-bnb4-sgmat.wgsl.jinja)
@@ -78,7 +78,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
78
  ## Use with `@huggingface/kernels`
79
 
80
  ```sh
81
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
82
  ```
83
 
84
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
68
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
69
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
70
  - [`test.json`](build/webgpu/test.json) — correctness cases
71
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
72
  - [`cast-scalar-x4.wgsl.jinja`](build/webgpu/cast-scalar-x4.wgsl.jinja)
73
  - [`matmul-bnb4-gemv.wgsl.jinja`](build/webgpu/matmul-bnb4-gemv.wgsl.jinja)
74
  - [`matmul-bnb4-sgmat.wgsl.jinja`](build/webgpu/matmul-bnb4-sgmat.wgsl.jinja)
 
78
  ## Use with `@huggingface/kernels`
79
 
80
  ```sh
81
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
82
  ```
83
 
84
  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
@@ -122,7 +122,7 @@
122
  {
123
  "name": "bnb4-nf4-speculative-m15-k4096-n4096-b64-band",
124
  "provenance": {
125
- "notes": "Fifteen rows represent a speculative decoder's verification width and exercise the tiled route immediately below M=16. `nf4-tiled-speculative-band-m16-control` brackets that boundary."
126
  },
127
  "preset": "model",
128
  "vars": { "M": 15, "K": 4096, "N": 4096, "blockSize": 64 },
 
122
  {
123
  "name": "bnb4-nf4-speculative-m15-k4096-n4096-b64-band",
124
  "provenance": {
125
+ "notes": "Measures a 15-row NF4 multiply (K=4096, N=4096, block size 64), the verification width of a speculative decoder. `bnb4-nf4-speculative-m16-k4096-n4096-b64-tile-floor-control` is the same multiply with 16 rows."
126
  },
127
  "preset": "model",
128
  "vars": { "M": 15, "K": 4096, "N": 4096, "blockSize": 64 },
build/webgpu/manifest.json CHANGED
@@ -72,6 +72,21 @@
72
  "sgmatWidenOutBytes": "dim(shapes.aT, 0) * attrs.N * 4",
73
  "sgmatWidenFits": "sgmatWidenBytes <= device.limits.maxStorageBufferBindingSize and sgmatWidenBytes <= device.limits.maxBufferSize and sgmatWidenOutBytes <= device.limits.maxStorageBufferBindingSize and sgmatWidenOutBytes <= device.limits.maxBufferSize",
74
  "quantType": "attrs.quant_type",
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
75
  "ABSMAX_LEN": "ceilDiv(attrs.N * attrs.K, attrs.block_size)",
76
  "stagedGeometryValid": "tunables.SGMAT_TILE_COLS >= sgmatSubCols and tunables.SGMAT_TILE_COLS % sgmatSubCols == 0 and tunables.SGMAT_TILE_K >= sgmatMatrixSize and tunables.SGMAT_TILE_K % sgmatMatrixSize == 0",
77
  "stagedRowSubtiles": "max(1, min(sgmatRowSubtiles, pow2ceil(ceilDiv(dim(shapes.aT, 0), sgmatMatrixSize)), floor(deviceWorkgroupCap / (32 * max(1, sgmatColSubtiles))), floor(tunables.SGMAT_TILE_K / sgmatMatrixSize), floor((device.limits.maxComputeWorkgroupStorageSize - sgmatWorkgroupStorageBytes) / (4 * sgmatMatrixSize * max(1, tunables.SGMAT_TILE_K)))))",
@@ -98,17 +113,11 @@
98
  "hybridResourcesFit": "prefixResourcesFit and hybridStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and hybridDispatchRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"
99
  },
100
  "bindings": {
101
- "a": { "arg": "aT", "buffer": "read-only-storage", "elementType": "$aScalar" },
102
- "b": { "arg": "bT", "buffer": "read-only-storage", "elementType": "u32" },
103
- "absmax": {
104
- "arg": "absmaxT",
105
- "buffer": "read-only-storage",
106
- "elementType": "$absmaxScalar",
107
- "length": "$ABSMAX_LEN"
108
- },
109
- "y": { "arg": "yT", "buffer": "storage", "elementType": "$aScalar" },
110
  "params": {
111
- "buffer": "uniform",
112
  "struct": [
113
  { "name": "rows", "type": "u32", "value": "dim(shapes.aT, 0)" },
114
  { "name": "K", "type": "u32", "value": "attrs.K" },
@@ -128,24 +137,9 @@
128
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
129
  },
130
  "derive": {
131
- "K": "attrs.K",
132
- "N": "attrs.N",
133
- "blockSize": "attrs.block_size",
134
  "tileRows": "prefixTileRows",
135
- "tileCols": "tunables.SGMAT_TILE_COLS",
136
- "tileK": "tunables.SGMAT_TILE_K",
137
  "subRows": "prefixTileRows / sgmatRowSubtiles",
138
- "subCols": "sgmatSubCols",
139
- "matrixSize": "sgmatMatrixSize",
140
  "rowMatrices": "prefixTileRows / (sgmatRowSubtiles * sgmatMatrixSize)",
141
- "colMatrices": "sgmatSubCols / sgmatMatrixSize",
142
- "loadWidth": "sgmatLoadWidth",
143
- "rowSubtiles": "sgmatRowSubtiles",
144
- "workgroupSize": "sgmatWorkgroupSize",
145
- "bLoadsPerRow": "sgmatBLoadsPerRow",
146
- "aScalar": "\"f32\"",
147
- "usesF16": false,
148
- "absmaxScalar": "\"f16\" if tensorDtypes.absmaxT == \"float16\" else \"f32\"",
149
  "hybridRows": true,
150
  "rowCount": "dim(shapes.aT, 0)",
151
  "fullDispatchRows": "sgmatFullRows / prefixTileRows",
@@ -170,28 +164,12 @@
170
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
171
  },
172
  "derive": {
173
- "K": "attrs.K",
174
- "N": "attrs.N",
175
- "blockSize": "attrs.block_size",
176
  "tileRows": "prefixTileRows",
177
- "tileCols": "tunables.SGMAT_TILE_COLS",
178
- "tileK": "tunables.SGMAT_TILE_K",
179
  "subRows": "prefixTileRows / sgmatRowSubtiles",
180
- "subCols": "sgmatSubCols",
181
- "matrixSize": "sgmatMatrixSize",
182
  "rowMatrices": "prefixTileRows / (sgmatRowSubtiles * sgmatMatrixSize)",
183
- "colMatrices": "sgmatSubCols / sgmatMatrixSize",
184
- "loadWidth": "sgmatLoadWidth",
185
- "rowSubtiles": "sgmatRowSubtiles",
186
- "workgroupSize": "sgmatWorkgroupSize",
187
- "bLoadsPerRow": "sgmatBLoadsPerRow",
188
  "aScalar": "\"f32\"",
189
- "usesF16": true,
190
- "absmaxScalar": "\"f16\" if tensorDtypes.absmaxT == \"float16\" else \"f32\"",
191
  "srcScalar": "\"f16\"",
192
  "outScalar": "\"f32\"",
193
- "wrapNarrowInt": false,
194
- "wrapSigned": false,
195
  "hybridRows": true,
196
  "rowCount": "dim(shapes.aT, 0)",
197
  "fullDispatchRows": "sgmatFullRows / prefixTileRows",
@@ -256,24 +234,13 @@
256
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
257
  },
258
  "derive": {
259
- "K": "attrs.K",
260
- "N": "attrs.N",
261
- "blockSize": "attrs.block_size",
262
  "tileRows": "stagedTileRows",
263
- "tileCols": "tunables.SGMAT_TILE_COLS",
264
- "tileK": "tunables.SGMAT_TILE_K",
265
  "subRows": "sgmatMatrixSize",
266
  "subCols": "stagedSubCols",
267
- "matrixSize": "sgmatMatrixSize",
268
  "rowMatrices": 1,
269
  "colMatrices": "stagedColumnMatrices",
270
- "loadWidth": "sgmatLoadWidth",
271
  "rowSubtiles": "stagedRowSubtiles",
272
  "workgroupSize": "stagedWorkgroupSize",
273
- "bLoadsPerRow": "sgmatBLoadsPerRow",
274
- "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
275
- "usesF16": "tensorDtypes.aT == \"float16\"",
276
- "absmaxScalar": "\"f16\" if tensorDtypes.absmaxT == \"float16\" else \"f32\"",
277
  "stageActivations": true,
278
  "rowCount": "dim(shapes.aT, 0)"
279
  },
@@ -290,12 +257,7 @@
290
  "id": "gemv",
291
  "priority": 20,
292
  "when": ["gemvShapeValid", "portableWorkgroupFits"],
293
- "derive": {
294
- "workgroupSize": "portableWorkgroupSize",
295
- "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
296
- "usesF16": "tensorDtypes.aT == \"float16\"",
297
- "absmaxScalar": "\"f16\" if tensorDtypes.absmaxT == \"float16\" else \"f32\""
298
- },
299
  "passes": [
300
  {
301
  "id": "main",
@@ -328,24 +290,9 @@
328
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
329
  },
330
  "derive": {
331
- "K": "attrs.K",
332
- "N": "attrs.N",
333
- "blockSize": "attrs.block_size",
334
  "tileRows": "sgmatTileRows",
335
- "tileCols": "tunables.SGMAT_TILE_COLS",
336
- "tileK": "tunables.SGMAT_TILE_K",
337
  "subRows": "sgmatSubRows",
338
- "subCols": "sgmatSubCols",
339
- "matrixSize": "sgmatMatrixSize",
340
- "rowMatrices": "sgmatSubRows / sgmatMatrixSize",
341
- "colMatrices": "sgmatSubCols / sgmatMatrixSize",
342
- "loadWidth": "sgmatLoadWidth",
343
- "rowSubtiles": "sgmatRowSubtiles",
344
- "workgroupSize": "sgmatWorkgroupSize",
345
- "bLoadsPerRow": "sgmatBLoadsPerRow",
346
- "aScalar": "\"f32\"",
347
- "usesF16": false,
348
- "absmaxScalar": "\"f16\" if tensorDtypes.absmaxT == \"float16\" else \"f32\""
349
  },
350
  "passes": [
351
  {
@@ -366,28 +313,12 @@
366
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
367
  },
368
  "derive": {
369
- "K": "attrs.K",
370
- "N": "attrs.N",
371
- "blockSize": "attrs.block_size",
372
  "tileRows": "sgmatTileRows",
373
- "tileCols": "tunables.SGMAT_TILE_COLS",
374
- "tileK": "tunables.SGMAT_TILE_K",
375
  "subRows": "sgmatSubRows",
376
- "subCols": "sgmatSubCols",
377
- "matrixSize": "sgmatMatrixSize",
378
  "rowMatrices": "sgmatSubRows / sgmatMatrixSize",
379
- "colMatrices": "sgmatSubCols / sgmatMatrixSize",
380
- "loadWidth": "sgmatLoadWidth",
381
- "rowSubtiles": "sgmatRowSubtiles",
382
- "workgroupSize": "sgmatWorkgroupSize",
383
- "bLoadsPerRow": "sgmatBLoadsPerRow",
384
  "aScalar": "\"f32\"",
385
- "usesF16": true,
386
- "absmaxScalar": "\"f16\" if tensorDtypes.absmaxT == \"float16\" else \"f32\"",
387
  "srcScalar": "\"f16\"",
388
- "outScalar": "\"f32\"",
389
- "wrapNarrowInt": false,
390
- "wrapSigned": false
391
  },
392
  "intermediates": [
393
  { "id": "aF32", "dtype": "float32", "shape": "[numel(shapes.aT)]" },
@@ -443,13 +374,7 @@
443
  "id": "tiled",
444
  "priority": 10,
445
  "when": ["tileEligible", "tileWorkgroupFits"],
446
- "derive": {
447
- "tileK": "tunables.PORTABLE_TILE_K",
448
- "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
449
- "usesF16": "tensorDtypes.aT == \"float16\"",
450
- "absmaxScalar": "\"f16\" if tensorDtypes.absmaxT == \"float16\" else \"f32\"",
451
- "tileRows": "portableTileRows"
452
- },
453
  "passes": [
454
  {
455
  "id": "main",
@@ -463,12 +388,7 @@
463
  "id": "scalar",
464
  "priority": 0,
465
  "when": ["commonShapeValid", "portableWorkgroupSize > 0"],
466
- "derive": {
467
- "workgroupSize": "portableWorkgroupSize",
468
- "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
469
- "usesF16": "tensorDtypes.aT == \"float16\"",
470
- "absmaxScalar": "\"f16\" if tensorDtypes.absmaxT == \"float16\" else \"f32\""
471
- },
472
  "passes": [
473
  {
474
  "id": "main",
 
72
  "sgmatWidenOutBytes": "dim(shapes.aT, 0) * attrs.N * 4",
73
  "sgmatWidenFits": "sgmatWidenBytes <= device.limits.maxStorageBufferBindingSize and sgmatWidenBytes <= device.limits.maxBufferSize and sgmatWidenOutBytes <= device.limits.maxStorageBufferBindingSize and sgmatWidenOutBytes <= device.limits.maxBufferSize",
74
  "quantType": "attrs.quant_type",
75
+ "K": "attrs.K",
76
+ "N": "attrs.N",
77
+ "blockSize": "attrs.block_size",
78
+ "tileCols": "tunables.SGMAT_TILE_COLS",
79
+ "tileK": "tunables.SGMAT_TILE_K",
80
+ "matrixSize": "sgmatMatrixSize",
81
+ "subCols": "sgmatSubCols",
82
+ "colMatrices": "sgmatSubCols / sgmatMatrixSize",
83
+ "loadWidth": "sgmatLoadWidth",
84
+ "rowSubtiles": "sgmatRowSubtiles",
85
+ "workgroupSize": "sgmatWorkgroupSize",
86
+ "bLoadsPerRow": "sgmatBLoadsPerRow",
87
+ "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
88
+ "usesF16": "tensorDtypes.aT == \"float16\"",
89
+ "absmaxScalar": "\"f16\" if tensorDtypes.absmaxT == \"float16\" else \"f32\"",
90
  "ABSMAX_LEN": "ceilDiv(attrs.N * attrs.K, attrs.block_size)",
91
  "stagedGeometryValid": "tunables.SGMAT_TILE_COLS >= sgmatSubCols and tunables.SGMAT_TILE_COLS % sgmatSubCols == 0 and tunables.SGMAT_TILE_K >= sgmatMatrixSize and tunables.SGMAT_TILE_K % sgmatMatrixSize == 0",
92
  "stagedRowSubtiles": "max(1, min(sgmatRowSubtiles, pow2ceil(ceilDiv(dim(shapes.aT, 0), sgmatMatrixSize)), floor(deviceWorkgroupCap / (32 * max(1, sgmatColSubtiles))), floor(tunables.SGMAT_TILE_K / sgmatMatrixSize), floor((device.limits.maxComputeWorkgroupStorageSize - sgmatWorkgroupStorageBytes) / (4 * sgmatMatrixSize * max(1, tunables.SGMAT_TILE_K)))))",
 
113
  "hybridResourcesFit": "prefixResourcesFit and hybridStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and hybridDispatchRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"
114
  },
115
  "bindings": {
116
+ "a": { "arg": "aT", "elementType": "$aScalar" },
117
+ "b": { "arg": "bT", "elementType": "u32" },
118
+ "absmax": { "arg": "absmaxT", "elementType": "$absmaxScalar", "length": "$ABSMAX_LEN" },
119
+ "y": { "arg": "yT", "elementType": "$aScalar" },
 
 
 
 
 
120
  "params": {
 
121
  "struct": [
122
  { "name": "rows", "type": "u32", "value": "dim(shapes.aT, 0)" },
123
  { "name": "K", "type": "u32", "value": "attrs.K" },
 
137
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
138
  },
139
  "derive": {
 
 
 
140
  "tileRows": "prefixTileRows",
 
 
141
  "subRows": "prefixTileRows / sgmatRowSubtiles",
 
 
142
  "rowMatrices": "prefixTileRows / (sgmatRowSubtiles * sgmatMatrixSize)",
 
 
 
 
 
 
 
 
143
  "hybridRows": true,
144
  "rowCount": "dim(shapes.aT, 0)",
145
  "fullDispatchRows": "sgmatFullRows / prefixTileRows",
 
164
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
165
  },
166
  "derive": {
 
 
 
167
  "tileRows": "prefixTileRows",
 
 
168
  "subRows": "prefixTileRows / sgmatRowSubtiles",
 
 
169
  "rowMatrices": "prefixTileRows / (sgmatRowSubtiles * sgmatMatrixSize)",
 
 
 
 
 
170
  "aScalar": "\"f32\"",
 
 
171
  "srcScalar": "\"f16\"",
172
  "outScalar": "\"f32\"",
 
 
173
  "hybridRows": true,
174
  "rowCount": "dim(shapes.aT, 0)",
175
  "fullDispatchRows": "sgmatFullRows / prefixTileRows",
 
234
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
235
  },
236
  "derive": {
 
 
 
237
  "tileRows": "stagedTileRows",
 
 
238
  "subRows": "sgmatMatrixSize",
239
  "subCols": "stagedSubCols",
 
240
  "rowMatrices": 1,
241
  "colMatrices": "stagedColumnMatrices",
 
242
  "rowSubtiles": "stagedRowSubtiles",
243
  "workgroupSize": "stagedWorkgroupSize",
 
 
 
 
244
  "stageActivations": true,
245
  "rowCount": "dim(shapes.aT, 0)"
246
  },
 
257
  "id": "gemv",
258
  "priority": 20,
259
  "when": ["gemvShapeValid", "portableWorkgroupFits"],
260
+ "derive": { "workgroupSize": "portableWorkgroupSize" },
 
 
 
 
 
261
  "passes": [
262
  {
263
  "id": "main",
 
290
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
291
  },
292
  "derive": {
 
 
 
293
  "tileRows": "sgmatTileRows",
 
 
294
  "subRows": "sgmatSubRows",
295
+ "rowMatrices": "sgmatSubRows / sgmatMatrixSize"
 
 
 
 
 
 
 
 
 
 
296
  },
297
  "passes": [
298
  {
 
313
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
314
  },
315
  "derive": {
 
 
 
316
  "tileRows": "sgmatTileRows",
 
 
317
  "subRows": "sgmatSubRows",
 
 
318
  "rowMatrices": "sgmatSubRows / sgmatMatrixSize",
 
 
 
 
 
319
  "aScalar": "\"f32\"",
 
 
320
  "srcScalar": "\"f16\"",
321
+ "outScalar": "\"f32\""
 
 
322
  },
323
  "intermediates": [
324
  { "id": "aF32", "dtype": "float32", "shape": "[numel(shapes.aT)]" },
 
374
  "id": "tiled",
375
  "priority": 10,
376
  "when": ["tileEligible", "tileWorkgroupFits"],
377
+ "derive": { "tileK": "tunables.PORTABLE_TILE_K", "tileRows": "portableTileRows" },
 
 
 
 
 
 
378
  "passes": [
379
  {
380
  "id": "main",
 
388
  "id": "scalar",
389
  "priority": 0,
390
  "when": ["commonShapeValid", "portableWorkgroupSize > 0"],
391
+ "derive": { "workgroupSize": "portableWorkgroupSize" },
 
 
 
 
 
392
  "passes": [
393
  {
394
  "id": "main",
build/webgpu/matmul-bnb4-gemv.wgsl.jinja CHANGED
@@ -21,14 +21,14 @@ var<private> quant_map: array<f32, 16> = array<f32, 16>(
21
  -0.0, -5.208333333e-03, -0.66666667, -1.0,
22
  -0.33333333, -0.5, -0.16666667, -0.25
23
  );
24
- {%- else %}
25
  var<private> quant_map: array<f32, 16> = array<f32, 16>(
26
  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
27
  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
28
  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
29
  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
30
  );
31
- {%- endif %}
32
 
33
  // One dequantized weight at flat index `flat` into the [N, K] weight.
34
  fn dequant_weight(flat: u32) -> f32 {
@@ -37,8 +37,6 @@ fn dequant_weight(flat: u32) -> f32 {
37
  return quant_map[code] * f32(absmax[flat / params.blockSize]);
38
  }
39
 
40
-
41
-
42
  // Decode (M == 1): one workgroup owns each output column, with lanes striding K.
43
  // Each lane keeps one accumulator and the workgroup closes with a tree reduction.
44
  @compute @workgroup_size(WG, 1, 1)
 
21
  -0.0, -5.208333333e-03, -0.66666667, -1.0,
22
  -0.33333333, -0.5, -0.16666667, -0.25
23
  );
24
+ {% else %}
25
  var<private> quant_map: array<f32, 16> = array<f32, 16>(
26
  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
27
  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
28
  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
29
  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
30
  );
31
+ {% endif %}
32
 
33
  // One dequantized weight at flat index `flat` into the [N, K] weight.
34
  fn dequant_weight(flat: u32) -> f32 {
 
37
  return quant_map[code] * f32(absmax[flat / params.blockSize]);
38
  }
39
 
 
 
40
  // Decode (M == 1): one workgroup owns each output column, with lanes striding K.
41
  // Each lane keeps one accumulator and the workgroup closes with a tree reduction.
42
  @compute @workgroup_size(WG, 1, 1)
build/webgpu/matmul-bnb4-sgmat.wgsl.jinja CHANGED
@@ -9,7 +9,6 @@
9
  {% if usesF16 %}
10
  // The widened path keeps the activations and the output in f32 for the matrix
11
  // units, but the quantization scales stay in the request's own dtype.
12
- enable f16;
13
  {% endif %}
14
  enable subgroups;
15
  {% if pinSubgroupSize32 %}
@@ -18,7 +17,6 @@ enable subgroup_size_control;
18
  enable chromium_experimental_subgroup_matrix;
19
  diagnostic(off, chromium.subgroup_matrix_uniformity);
20
 
21
-
22
  {{ env.wgsl.resourceDeclarations }}
23
 
24
  const K: u32 = {{ K }}u;
@@ -59,17 +57,14 @@ var<private> quant_map: array<f32, 16> = array<f32, 16>(
59
  -0.0, -5.208333333e-03, -0.66666667, -1.0,
60
  -0.33333333, -0.5, -0.16666667, -0.25
61
  );
62
- {%- else %}
63
  var<private> quant_map: array<f32, 16> = array<f32, 16>(
64
  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
65
  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
66
  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
67
  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
68
  );
69
- {%- endif %}
70
-
71
-
72
-
73
 
74
  fn load_b(tile_base: u32, k_base: u32, row: u32, col_group: u32) {
75
  let col = col_group * LOAD_WIDTH;
 
9
  {% if usesF16 %}
10
  // The widened path keeps the activations and the output in f32 for the matrix
11
  // units, but the quantization scales stay in the request's own dtype.
 
12
  {% endif %}
13
  enable subgroups;
14
  {% if pinSubgroupSize32 %}
 
17
  enable chromium_experimental_subgroup_matrix;
18
  diagnostic(off, chromium.subgroup_matrix_uniformity);
19
 
 
20
  {{ env.wgsl.resourceDeclarations }}
21
 
22
  const K: u32 = {{ K }}u;
 
57
  -0.0, -5.208333333e-03, -0.66666667, -1.0,
58
  -0.33333333, -0.5, -0.16666667, -0.25
59
  );
60
+ {% else %}
61
  var<private> quant_map: array<f32, 16> = array<f32, 16>(
62
  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
63
  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
64
  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
65
  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
66
  );
67
+ {% endif %}
 
 
 
68
 
69
  fn load_b(tile_base: u32, k_base: u32, row: u32, col_group: u32) {
70
  let col = col_group * LOAD_WIDTH;
build/webgpu/matmul-bnb4-tiled.wgsl.jinja CHANGED
@@ -41,15 +41,14 @@ var<private> quant_map: array<f32, 16> = array<f32, 16>(
41
  -0.0, -5.208333333e-03, -0.66666667, -1.0,
42
  -0.33333333, -0.5, -0.16666667, -0.25
43
  );
44
- {%- else %}
45
  var<private> quant_map: array<f32, 16> = array<f32, 16>(
46
  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
47
  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
48
  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
49
  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
50
  );
51
- {%- endif %}
52
-
53
 
54
  // Four consecutive dequantized weights starting at flat index `flat` into the
55
  // [N, K] weight. Consecutive indices share a stored byte, so the run costs two
@@ -82,7 +81,6 @@ fn dequant_weight4(flat: u32) -> vec4<f32> {
82
  return vec4<f32>(quant_map[codes.x], quant_map[codes.y], quant_map[codes.z], quant_map[codes.w]) * scales;
83
  }
84
 
85
-
86
  @compute @workgroup_size(16, 8, 1)
87
  fn main(
88
  @builtin(workgroup_id) wg: vec3<u32>,
 
41
  -0.0, -5.208333333e-03, -0.66666667, -1.0,
42
  -0.33333333, -0.5, -0.16666667, -0.25
43
  );
44
+ {% else %}
45
  var<private> quant_map: array<f32, 16> = array<f32, 16>(
46
  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
47
  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
48
  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
49
  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
50
  );
51
+ {% endif %}
 
52
 
53
  // Four consecutive dequantized weights starting at flat index `flat` into the
54
  // [N, K] weight. Consecutive indices share a stored byte, so the run costs two
 
81
  return vec4<f32>(quant_map[codes.x], quant_map[codes.y], quant_map[codes.z], quant_map[codes.w]) * scales;
82
  }
83
 
 
84
  @compute @workgroup_size(16, 8, 1)
85
  fn main(
86
  @builtin(workgroup_id) wg: vec3<u32>,
build/webgpu/matmul-bnb4.wgsl.jinja CHANGED
@@ -1,3 +1,8 @@
 
 
 
 
 
1
  {{ env.wgsl.resourceDeclarations }}
2
 
3
  const WG: u32 = {{ workgroupSize }}u;
@@ -19,14 +24,14 @@ var<private> quant_map: array<f32, 16> = array<f32, 16>(
19
  -0.0, -5.208333333e-03, -0.66666667, -1.0,
20
  -0.33333333, -0.5, -0.16666667, -0.25
21
  );
22
- {%- else %}
23
  var<private> quant_map: array<f32, 16> = array<f32, 16>(
24
  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
25
  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
26
  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
27
  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
28
  );
29
- {%- endif %}
30
 
31
  // One dequantized weight at flat index `flat` into the [N, K] weight.
32
  fn dequant_weight(flat: u32) -> f32 {
@@ -35,14 +40,9 @@ fn dequant_weight(flat: u32) -> f32 {
35
  return quant_map[code] * f32(absmax[flat / params.blockSize]);
36
  }
37
 
38
-
39
-
40
  @compute @workgroup_size(WG, 1, 1)
41
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
42
- // 2D-folded flat output-element index: gid.y carries the high bits past the
43
- // per-axis dispatch fold width. With no fold this
44
- // reduces to gid.x; the index >= total guard drops the tail.
45
- let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
46
  let total = params.rows * params.N;
47
 
48
  if (index >= total) {
 
1
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
2
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
3
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
+ // per-axis workgroup fold width.
5
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
6
  {{ env.wgsl.resourceDeclarations }}
7
 
8
  const WG: u32 = {{ workgroupSize }}u;
 
24
  -0.0, -5.208333333e-03, -0.66666667, -1.0,
25
  -0.33333333, -0.5, -0.16666667, -0.25
26
  );
27
+ {% else %}
28
  var<private> quant_map: array<f32, 16> = array<f32, 16>(
29
  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
30
  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
31
  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
32
  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
33
  );
34
+ {% endif %}
35
 
36
  // One dequantized weight at flat index `flat` into the [N, K] weight.
37
  fn dequant_weight(flat: u32) -> f32 {
 
40
  return quant_map[code] * f32(absmax[flat / params.blockSize]);
41
  }
42
 
 
 
43
  @compute @workgroup_size(WG, 1, 1)
44
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
45
+ {{ flat_index_2d("WG", "index", "") }}
 
 
 
46
  let total = params.rows * params.N;
47
 
48
  if (index >= total) {
build/webgpu/metadata.json CHANGED
@@ -1,25 +1,25 @@
1
  {
2
  "name": "com.microsoft.MatMulBnb4",
3
- "id": "_com_microsoft_matmulbnb4_webgpu_237048a",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "/Z1zFaT4D+0TEFMxG/zqun5CjI9irI93GNGwLZdLXHQ=",
11
  "cast-scalar-x4.wgsl.jinja": "s8Y9E2+D7cFgYTUEJhiuPsUzvHjv2mXVoiRmv0A0vTU=",
12
- "manifest.json": "TEJfUwu6Isit+t0fBCRp0Ny+uC8hUjSe2jgG5kCFcEg=",
13
- "matmul-bnb4-gemv.wgsl.jinja": "7W3L5RH+5dk5/dKNnjwVFtBt+xgCxorJJ04/8zHrySw=",
14
- "matmul-bnb4-sgmat.wgsl.jinja": "9WS5F6zd/KFztTpbZjBgJe7WIl0l7uiRelcDv/9YbBc=",
15
- "matmul-bnb4-tiled.wgsl.jinja": "bADcd5OFuhrZ33Qntqh1PXREZolJZ+RKS+VlweAjPo0=",
16
- "matmul-bnb4.wgsl.jinja": "6rtf2K0op1IJv2+l66HJjq8XnBATGyefxit0Hj5QREI=",
17
  "test.json": "KPAf0hjucty64auyD+xOS5RAGDEPYeAlUSJPhbig7CA="
18
  }
19
  },
20
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
21
  "webgpu": {
22
- "manifestSpec": "2.0",
23
  "variants": {
24
  "sgmat_hybrid_rows": ["matmul-bnb4-sgmat.wgsl.jinja"],
25
  "sgmat_hybrid_rows_widened": ["cast-scalar-x4.wgsl.jinja", "matmul-bnb4-sgmat.wgsl.jinja"],
 
1
  {
2
  "name": "com.microsoft.MatMulBnb4",
3
+ "id": "_com_microsoft_matmulbnb4_webgpu_2efe1ae",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "Ixtu6FBToeEWEbpbM/+lGVS3jt0IPtrmZrxRfxo7Uj8=",
11
  "cast-scalar-x4.wgsl.jinja": "s8Y9E2+D7cFgYTUEJhiuPsUzvHjv2mXVoiRmv0A0vTU=",
12
+ "manifest.json": "TTpD+bhv4VB93QlNQ4Z2xGHvG4ClFxnokPMIbKMH2Bc=",
13
+ "matmul-bnb4-gemv.wgsl.jinja": "5mV78S/5vyeTFu8anlXAsKA+OVz7iy3fmo990lgf/yU=",
14
+ "matmul-bnb4-sgmat.wgsl.jinja": "A0jgHcIQrjdodubOI3bgHSYU9SsGew+/3IJwkMMr6eo=",
15
+ "matmul-bnb4-tiled.wgsl.jinja": "auoAXKsuYOXVbPmzjXMm6IQ5L1HxzjDWaPTSxD+Wd/s=",
16
+ "matmul-bnb4.wgsl.jinja": "81lKMMKBPnPeNbcFXiZZe4KBGefTeBfNaHH1LSgP/6U=",
17
  "test.json": "KPAf0hjucty64auyD+xOS5RAGDEPYeAlUSJPhbig7CA="
18
  }
19
  },
20
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
21
  "webgpu": {
22
+ "manifestSpec": "2.1",
23
  "variants": {
24
  "sgmat_hybrid_rows": ["matmul-bnb4-sgmat.wgsl.jinja"],
25
  "sgmat_hybrid_rows_widened": ["cast-scalar-x4.wgsl.jinja", "matmul-bnb4-sgmat.wgsl.jinja"],