| { |
| "domain": "com.microsoft", |
| "name": "MatMulBnb4", |
| "sinceVersion": 1, |
| "description": "Computes `A @ dequant(B)^T` where `B` uses bitsandbytes 4-bit quantization: `quant_type = 0` selects FP4 and `quant_type = 1` selects NF4. Supports rank-2 float16/float32 `A`, `transB = 1`, and `training_mode = 0`; rank-1 and rank-3-or-higher `A`, bfloat16, `transB = 0`, and training are not implemented. `B` is the flattened `[N, K]` weight, two codes per byte with the even flat index in the high nibble. Each code indexes a fixed 16-entry codebook, and the value is `codebook[code] * absmax[flat_index / block_size]`.", |
| "inputs": [ |
| { "role": "A", "dtype": "T1", "rank": 2, "description": "Float input matrix of shape `(M, K)`, not quantized." }, |
| { |
| "role": "B", |
| "dtype": "T2", |
| "rank": 1, |
| "description": "The `[N, K]` weight, flattened and quantized to 4 bits, stored as `(N * K + 1) / 2` bytes; the ONNX type is uint8 (this WebGPU implementation reads one widened u32 per stored byte)." |
| }, |
| { |
| "role": "absmax", |
| "dtype": "T1", |
| "rank": 1, |
| "description": "Per-block absolute-maximum dequantization scales of shape `((N * K + block_size - 1) / block_size)`, same dtype as A." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "Y", |
| "dtype": "T1", |
| "rank": 2, |
| "shape": "[dim(shapes.A, 0), attrs.N]", |
| "description": "Result of `A` multiplied by the dequantized, transposed weight matrix, with shape `(M, N)` and the same dtype as `A`." |
| } |
| ], |
| "attributes": { "training_mode": 0, "transB": 1 }, |
| "attributeConstraints": { |
| "K": { "required": true }, |
| "N": { "required": true }, |
| "block_size": { "required": true }, |
| "quant_type": { "required": true, "values": [0, 1] }, |
| "training_mode": { "values": [0] }, |
| "transB": { "values": [1] } |
| }, |
| "attributeDescriptions": { |
| "K": "Input feature count (the shared dimension).", |
| "N": "Output feature count.", |
| "block_size": "Number of weights sharing one absmax scale; a power of two, at least 16.", |
| "quant_type": "Codebook selector: 0 = FP4, 1 = NF4.", |
| "training_mode": "Whether training outputs are requested. This inference-only implementation supports the standard default value 0.", |
| "transB": "Whether the quantized weight is stored transposed. This implementation supports the standard default value 1." |
| }, |
| "typeConstraints": { "T1": ["float32", "float16"], "T2": ["uint8"] }, |
| "args": { |
| "aT": { "kind": "tensor", "semantic": "A", "role": "input" }, |
| "bT": { "kind": "tensor", "semantic": "B", "role": "input" }, |
| "absmaxT": { "kind": "tensor", "semantic": "absmax", "role": "input" }, |
| "yT": { "kind": "tensor", "semantic": "Y", "role": "output" } |
| }, |
| "tunables": { |
| "WORKGROUP_SIZE": 64, |
| "TILE_MIN_M": 16, |
| "PORTABLE_TILE_K": 16, |
| "SGMAT_TILE_ROWS": 64, |
| "SGMAT_TALL_TILE_ROWS": 128, |
| "SGMAT_TALL_MIN_M": 128, |
| "SGMAT_TILE_COLS": 64, |
| "SGMAT_TILE_K": 32 |
| }, |
| "bindingSets": { |
| "main": [ |
| { |
| "name": "a", |
| "arg": "aT", |
| "semantic": "A", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { "name": "b", "arg": "bT", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "absmax", |
| "arg": "absmaxT", |
| "semantic": "absmax", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$absmaxScalar", |
| "length": "$ABSMAX_LEN" |
| }, |
| { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "rows", "type": "u32", "value": "dim(shapes.A, 0)" }, |
| { "name": "K", "type": "u32", "value": "attrs.K" }, |
| { "name": "N", "type": "u32", "value": "attrs.N" }, |
| { "name": "blockSize", "type": "u32", "value": "attrs.block_size" } |
| ] |
| } |
| } |
| ], |
| "gemv": [ |
| { |
| "name": "a", |
| "arg": "aT", |
| "semantic": "A", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { "name": "b", "arg": "bT", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "absmax", |
| "arg": "absmaxT", |
| "semantic": "absmax", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$absmaxScalar", |
| "length": "$ABSMAX_LEN" |
| }, |
| { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "K", "type": "u32", "value": "attrs.K" }, |
| { "name": "N", "type": "u32", "value": "attrs.N" }, |
| { "name": "blockSize", "type": "u32", "value": "attrs.block_size" } |
| ] |
| } |
| } |
| ], |
| "sgmat": [ |
| { |
| "name": "a", |
| "arg": "aT", |
| "semantic": "A", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { "name": "b", "arg": "bT", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "absmax", |
| "arg": "absmaxT", |
| "semantic": "absmax", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$absmaxScalar", |
| "length": "$ABSMAX_LEN" |
| }, |
| { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" } |
| ], |
| "castA": [ |
| { |
| "name": "x", |
| "arg": "aT", |
| "semantic": "A", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$srcScalar" |
| }, |
| { "name": "y", "semantic": "aF32", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.A)" }] } |
| } |
| ], |
| "sgmatWiden": [ |
| { "name": "a", "semantic": "aF32", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "b", "arg": "bT", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "absmax", |
| "arg": "absmaxT", |
| "semantic": "absmax", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$absmaxScalar", |
| "length": "$ABSMAX_LEN" |
| }, |
| { "name": "y", "semantic": "yF32", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "castY": [ |
| { "name": "x", "semantic": "yF32", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$srcScalar" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [{ "name": "count", "type": "u32", "value": "dim(shapes.A, 0) * attrs.N" }] |
| } |
| } |
| ] |
| }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32", |
| "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", |
| "packedBytesExpected": "ceilDiv(attrs.N * attrs.K, 2)", |
| "absmaxCountExpected": "ceilDiv(attrs.N * attrs.K, attrs.block_size)", |
| "aFloatOk": "(tensorDtypes.A == \"float32\" or tensorDtypes.A == \"float16\") and f16Ok(tensorDtypes.A)", |
| "portableWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "commonShapeValid": "ranks.A == 2 and ranks.B == 1 and ranks.absmax == 1 and ranks.Y == 2 and aFloatOk and tensorDtypes.B == \"uint8\" and tensorDtypes.absmax == tensorDtypes.A and tensorDtypes.Y == tensorDtypes.A and attrs.K > 0 and attrs.N > 0 and attrs.block_size >= 16 and pow2ceil(attrs.block_size) == attrs.block_size and dim(shapes.A, 1) == attrs.K and dim(shapes.B, 0) == packedBytesExpected and dim(shapes.absmax, 0) == absmaxCountExpected and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == attrs.N", |
| "gemvShapeValid": "commonShapeValid and dim(shapes.A, 0) == 1", |
| "portableWorkgroupFits": "portableWorkgroupSize > 0 and portableWorkgroupSize * 16 <= device.limits.maxComputeWorkgroupStorageSize", |
| "portableTileKValid": "tunables.PORTABLE_TILE_K >= 8 and tunables.PORTABLE_TILE_K % 8 == 0", |
| "tileEligible": "commonShapeValid and dim(shapes.A, 0) >= tunables.TILE_MIN_M and portableTileKValid", |
| "tileWorkgroupStorageBytes": "8 * 64 * tunables.PORTABLE_TILE_K", |
| "tileWorkgroupFits": "16 <= device.limits.maxComputeWorkgroupSizeX and 8 <= device.limits.maxComputeWorkgroupSizeY and 128 <= device.limits.maxComputeInvocationsPerWorkgroup and tileWorkgroupStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(attrs.N, 64) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(dim(shapes.A, 0), 64) <= device.limits.maxComputeWorkgroupsPerDimension", |
| "sgmatMatrixSize": "8", |
| "sgmatTileRows": "tunables.SGMAT_TALL_TILE_ROWS if dim(shapes.A, 0) >= tunables.SGMAT_TALL_MIN_M else tunables.SGMAT_TILE_ROWS", |
| "sgmatRowSubtiles": "4", |
| "sgmatSubRows": "sgmatTileRows / sgmatRowSubtiles", |
| "sgmatSubCols": "4 * sgmatMatrixSize", |
| "sgmatLoadWidth": "sgmatMatrixSize", |
| "sgmatColSubtiles": "tunables.SGMAT_TILE_COLS / sgmatSubCols", |
| "sgmatNumSubgroups": "sgmatRowSubtiles * sgmatColSubtiles", |
| "sgmatSubgroupSize": "device.adapterInfo.subgroupMinSize if has(device.adapterInfo, \"subgroupMinSize\") else 1", |
| "sgmatWorkgroupSize": "sgmatNumSubgroups * sgmatSubgroupSize", |
| "sgmatBLoadsPerRow": "tunables.SGMAT_TILE_K / sgmatLoadWidth", |
| "sgmatWorkgroupStorageBytes": "4 * tunables.SGMAT_TILE_COLS * tunables.SGMAT_TILE_K", |
| "sgmatDispatchN": "ceilDiv(attrs.N, tunables.SGMAT_TILE_COLS)", |
| "sgmatDispatchM": "ceilDiv(dim(shapes.A, 0), sgmatTileRows)", |
| "sgmatWorkgroupFits": "sgmatWorkgroupSize <= deviceWorkgroupCap and sgmatWorkgroupStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and sgmatDispatchN <= device.limits.maxComputeWorkgroupsPerDimension and sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", |
| "sgmatWidenBytes": "numel(shapes.A) * 4", |
| "sgmatWidenOutBytes": "dim(shapes.A, 0) * attrs.N * 4", |
| "sgmatWidenFits": "sgmatWidenBytes <= device.limits.maxStorageBufferBindingSize and sgmatWidenBytes <= device.limits.maxBufferSize and sgmatWidenOutBytes <= device.limits.maxStorageBufferBindingSize and sgmatWidenOutBytes <= device.limits.maxBufferSize" |
| }, |
| "constants": { "quantType": "attrs.quant_type", "ABSMAX_LEN": "ceilDiv(attrs.N * attrs.K, attrs.block_size)" }, |
| "variants": [ |
| { |
| "id": "gemv", |
| "priority": 20, |
| "when": ["gemvShapeValid", "portableWorkgroupFits"], |
| "constants": { |
| "workgroupSize": "portableWorkgroupSize", |
| "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", |
| "usesF16": "tensorDtypes.A == \"float16\"", |
| "absmaxScalar": "\"f16\" if tensorDtypes.absmax == \"float16\" else \"f32\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "matmul-bnb4-gemv.wgsl.jinja", |
| "bindings": "gemv", |
| "dispatch": { "workgroups": "attrs.N" } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat", |
| "priority": 15, |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "limits": { "maxComputeWorkgroupStorageSize": 8192 }, |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "when": ["commonShapeValid", "tensorDtypes.A == \"float32\"", "dim(shapes.A, 0) >= sgmatTileRows", "dim(shapes.A, 0) % sgmatTileRows == 0", "attrs.N % tunables.SGMAT_TILE_COLS == 0", "attrs.K % tunables.SGMAT_TILE_K == 0", "attrs.K % attrs.block_size == 0", "attrs.block_size % sgmatLoadWidth == 0", "wave32Effective", "sgmatWorkgroupFits"], |
| "constants": { |
| "K": "attrs.K", |
| "N": "attrs.N", |
| "blockSize": "attrs.block_size", |
| "tileRows": "sgmatTileRows", |
| "tileCols": "tunables.SGMAT_TILE_COLS", |
| "tileK": "tunables.SGMAT_TILE_K", |
| "subRows": "sgmatSubRows", |
| "subCols": "sgmatSubCols", |
| "matrixSize": "sgmatMatrixSize", |
| "rowMatrices": "sgmatSubRows / sgmatMatrixSize", |
| "colMatrices": "sgmatSubCols / sgmatMatrixSize", |
| "loadWidth": "sgmatLoadWidth", |
| "rowSubtiles": "sgmatRowSubtiles", |
| "workgroupSize": "sgmatWorkgroupSize", |
| "bLoadsPerRow": "sgmatBLoadsPerRow", |
| "aScalar": "\"f32\"", |
| "usesF16": false, |
| "absmaxScalar": "\"f16\" if tensorDtypes.absmax == \"float16\" else \"f32\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "matmul-bnb4-sgmat.wgsl.jinja", |
| "bindings": "sgmat", |
| "dispatch": { "x": "sgmatDispatchN", "y": "sgmatDispatchM" } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_widened", |
| "priority": 16, |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "limits": { "maxComputeWorkgroupStorageSize": 8192 }, |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "when": ["commonShapeValid", "tensorDtypes.A == \"float16\"", "dim(shapes.A, 0) >= sgmatTileRows", "dim(shapes.A, 0) % sgmatTileRows == 0", "attrs.N % tunables.SGMAT_TILE_COLS == 0", "attrs.K % tunables.SGMAT_TILE_K == 0", "attrs.K % attrs.block_size == 0", "attrs.block_size % sgmatLoadWidth == 0", "wave32Effective", "sgmatWorkgroupFits", "sgmatWidenFits"], |
| "constants": { |
| "K": "attrs.K", |
| "N": "attrs.N", |
| "blockSize": "attrs.block_size", |
| "tileRows": "sgmatTileRows", |
| "tileCols": "tunables.SGMAT_TILE_COLS", |
| "tileK": "tunables.SGMAT_TILE_K", |
| "subRows": "sgmatSubRows", |
| "subCols": "sgmatSubCols", |
| "matrixSize": "sgmatMatrixSize", |
| "rowMatrices": "sgmatSubRows / sgmatMatrixSize", |
| "colMatrices": "sgmatSubCols / sgmatMatrixSize", |
| "loadWidth": "sgmatLoadWidth", |
| "rowSubtiles": "sgmatRowSubtiles", |
| "workgroupSize": "sgmatWorkgroupSize", |
| "bLoadsPerRow": "sgmatBLoadsPerRow", |
| "aScalar": "\"f32\"", |
| "usesF16": true, |
| "absmaxScalar": "\"f16\" if tensorDtypes.absmax == \"float16\" else \"f32\"", |
| "srcScalar": "\"f16\"", |
| "outScalar": "\"f32\"", |
| "wrapNarrowInt": false, |
| "wrapSigned": false |
| }, |
| "passes": [ |
| { |
| "id": "widen_a", |
| "name": "MatMulBnb4.WidenActivations", |
| "shader": "cast-scalar-x4.wgsl.jinja", |
| "bindings": "castA", |
| "dispatch": { "threads": "ceilDiv(numel(shapes.A), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulBnb4.SubgroupMatrixWidened", |
| "shader": "matmul-bnb4-sgmat.wgsl.jinja", |
| "bindings": "sgmatWiden", |
| "dispatch": { "x": "sgmatDispatchN", "y": "sgmatDispatchM" } |
| }, |
| { |
| "id": "narrow_y", |
| "name": "MatMulBnb4.NarrowOutput", |
| "shader": "cast-scalar-x4.wgsl.jinja", |
| "bindings": "castY", |
| "constants": { "outScalar": "\"f16\"" }, |
| "dispatch": { |
| "threads": "ceilDiv(dim(shapes.A, 0) * attrs.N, 4)", |
| "workgroupSize": "tunables.WORKGROUP_SIZE" |
| } |
| } |
| ], |
| "description": "Runs the subgroup-matrix tier for a float16 request by widening the activations to float32 either side of the multiply. The matrix units this operator uses accumulate in float32 from float32 operands; the device's float16 configuration returns a float16 result, so feeding them float16 directly would drop the accumulator's precision.", |
| "intermediates": [ |
| { "id": "aF32", "dtype": "float32", "shape": "[numel(shapes.A)]" }, |
| { "id": "yF32", "dtype": "float32", "shape": "[dim(shapes.A, 0) * attrs.N]" } |
| ] |
| }, |
| { |
| "id": "tiled", |
| "priority": 10, |
| "when": ["tileEligible", "tileWorkgroupFits"], |
| "constants": { |
| "tileK": "tunables.PORTABLE_TILE_K", |
| "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", |
| "usesF16": "tensorDtypes.A == \"float16\"", |
| "absmaxScalar": "\"f16\" if tensorDtypes.absmax == \"float16\" else \"f32\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "matmul-bnb4-tiled.wgsl.jinja", |
| "bindings": "main", |
| "dispatch": { "x": "ceilDiv(attrs.N, 64)", "y": "ceilDiv(dim(shapes.A, 0), 64)" } |
| } |
| ] |
| }, |
| { |
| "id": "scalar", |
| "priority": 0, |
| "when": ["commonShapeValid", "portableWorkgroupSize > 0"], |
| "constants": { |
| "workgroupSize": "portableWorkgroupSize", |
| "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", |
| "usesF16": "tensorDtypes.A == \"float16\"", |
| "absmaxScalar": "\"f16\" if tensorDtypes.absmax == \"float16\" else \"f32\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "matmul-bnb4.wgsl.jinja", |
| "bindings": "main", |
| "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "constants.workgroupSize" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
