| { |
| "domain": "com.microsoft", |
| "name": "GatherBlockQuantized", |
| "sinceVersion": 1, |
| "description": "Gathers rows from a block-wise quantized weight matrix and dequantizes them. This inference implementation supports the standard `gather_axis = 0`, `quantize_axis = 1` matrix subset with uint8 `data`, 4-bit packed or 8-bit values, rank-1 non-negative in-bounds int64 `indices` projected to uint32 WebGPU storage, and float32 scales/output. Higher-rank gathers, negative indices, int32 indices, int4/uint4 data, 2-bit data, float16/bfloat16 output, and non-default axes are not implemented.", |
| "inputs": [ |
| { |
| "role": "data", |
| "dtype": "T1", |
| "rank": 2, |
| "description": "Constant uint8 weight matrix. With `bits = 4`, each byte stores two values low-nibble first; with `bits = 8`, each byte stores one value." |
| }, |
| { |
| "role": "indices", |
| "dtype": "Tind", |
| "rank": 1, |
| "description": "Non-negative logical int64 indices selecting rows from axis 0 of `data`. Every index must be less than the row count; values use checked uint32 WebGPU storage." |
| }, |
| { |
| "role": "scales", |
| "dtype": "T2", |
| "rank": 2, |
| "description": "Per-block dequantization scale factors of shape `(rows, ceil(output_columns / block_size))`." |
| }, |
| { |
| "role": "zero_points", |
| "dtype": "T1", |
| "rank": 2, |
| "optional": true, |
| "description": "Optional uint8 zero points. At 4 bits two zero points are packed per byte along the quantized axis, low-nibble first; at 8 bits the shape matches `scales`. If absent, uint8 data uses 2^(bits-1)." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "output", |
| "dtype": "T2", |
| "rank": 2, |
| "shape": "[dim(shapes.indices, 0), dim(shapes.data, 1) * (8 / attrs.bits)]", |
| "description": "Dequantized floating-point output rows corresponding to the gathered indices." |
| } |
| ], |
| "attributes": { "bits": 4, "block_size": 128, "gather_axis": 0, "quantize_axis": 1 }, |
| "attributeConstraints": { |
| "bits": { "values": [4, 8] }, |
| "gather_axis": { "values": [0] }, |
| "quantize_axis": { "values": [1] } |
| }, |
| "attributeDescriptions": { |
| "bits": "Bits per quantized value. The schema default is 4; this implementation supports 4 or 8.", |
| "block_size": "Number of values sharing a scale. Defaults to 128 and must be a power of two at least 16.", |
| "gather_axis": "Axis from which values are gathered. This matrix implementation supports the standard default, axis 0.", |
| "quantize_axis": "Axis split into quantization blocks. This matrix implementation supports the standard default, axis 1." |
| }, |
| "typeConstraints": { "T1": ["uint8"], "T2": ["float32"], "Tind": ["int64"] }, |
| "args": { |
| "dataT": { "kind": "tensor", "semantic": "data", "role": "input" }, |
| "indicesT": { "kind": "tensor", "semantic": "indices", "role": "input", "dtype": "uint32", "narrowing": "checked" }, |
| "scalesT": { "kind": "tensor", "semantic": "scales", "role": "input" }, |
| "zeroPointsT": { "kind": "tensor", "semantic": "zero_points", "role": "input", "required": false }, |
| "outputT": { "kind": "tensor", "semantic": "output", "role": "output" } |
| }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "bits": "attrs.bits", |
| "blockSize": "attrs.block_size", |
| "blockSizeOk": "blockSize >= 16 and pow2ceil(blockSize) == blockSize", |
| "outBlocks": "ceilDiv(dim(shapes.output, 1), blockSize)", |
| "zeroPointCols": "ceilDiv(outBlocks, 2) if bits == 4 else outBlocks", |
| "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "commonShapeValid": "blockSizeOk and ranks.data == 2 and ranks.indices == 1 and ranks.scales == 2 and ranks.output == 2 and tensorDtypes.indices == \"uint32\" and tensorDtypes.scales == \"float32\" and tensorDtypes.output == \"float32\" and dim(shapes.output, 0) == dim(shapes.indices, 0) and dim(shapes.scales, 0) == dim(shapes.data, 0) and dim(shapes.scales, 1) == outBlocks", |
| "q4ShapeValid": "commonShapeValid and tensorDtypes.data == \"uint8\" and dim(shapes.output, 1) == dim(shapes.data, 1) * 2", |
| "q8ShapeValid": "commonShapeValid and tensorDtypes.data == \"uint8\" and dim(shapes.output, 1) == dim(shapes.data, 1)", |
| "zeroPointsValid": "present.zeroPointsT and ranks.zero_points == 2 and tensorDtypes.zero_points == \"uint8\" and dim(shapes.zero_points, 0) == dim(shapes.data, 0) and dim(shapes.zero_points, 1) == zeroPointCols", |
| "noZeroMode": "not present.zeroPointsT", |
| "zeroMode": "zeroPointsValid", |
| "workgroupFits": "workgroupSize > 0", |
| "foldedDispatchFits": "ceil(ceil(numel(shapes.output) / device.limits.maxComputeWorkgroupsPerDimension) / workgroupSize) <= device.limits.maxComputeWorkgroupsPerDimension" |
| }, |
| "tunables": { "WORKGROUP_SIZE": 64 }, |
| "bindingSets": { |
| "noZero": [ |
| { |
| "name": "data", |
| "arg": "dataT", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$dataElement" |
| }, |
| { |
| "name": "indices", |
| "arg": "indicesT", |
| "semantic": "indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$indexScalar" |
| }, |
| { |
| "name": "scales", |
| "arg": "scalesT", |
| "semantic": "scales", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scaleScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputElement" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" }, |
| { "name": "cols", "type": "u32", "value": "dim(shapes.output, 1)" }, |
| { "name": "packedCols", "type": "u32", "value": "dim(shapes.data, 1)" }, |
| { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" }, |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" } |
| ] |
| } |
| } |
| ], |
| "zero": [ |
| { |
| "name": "data", |
| "arg": "dataT", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$dataElement" |
| }, |
| { |
| "name": "indices", |
| "arg": "indicesT", |
| "semantic": "indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$indexScalar" |
| }, |
| { |
| "name": "scales", |
| "arg": "scalesT", |
| "semantic": "scales", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scaleScalar" |
| }, |
| { |
| "name": "zero_points", |
| "arg": "zeroPointsT", |
| "semantic": "zero_points", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$zeroPointElement" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputElement" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" }, |
| { "name": "cols", "type": "u32", "value": "dim(shapes.output, 1)" }, |
| { "name": "packedCols", "type": "u32", "value": "dim(shapes.data, 1)" }, |
| { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" }, |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" } |
| ] |
| } |
| } |
| ], |
| "noZeroIo": [ |
| { |
| "name": "data", |
| "arg": "dataT", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$dataElement" |
| }, |
| { |
| "name": "indices", |
| "arg": "indicesT", |
| "semantic": "indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$indexScalar" |
| }, |
| { |
| "name": "scales", |
| "arg": "scalesT", |
| "semantic": "scales", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scaleScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputElement" |
| } |
| ], |
| "zeroIo": [ |
| { |
| "name": "data", |
| "arg": "dataT", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$dataElement" |
| }, |
| { |
| "name": "indices", |
| "arg": "indicesT", |
| "semantic": "indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$indexScalar" |
| }, |
| { |
| "name": "scales", |
| "arg": "scalesT", |
| "semantic": "scales", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scaleScalar" |
| }, |
| { |
| "name": "zero_points", |
| "arg": "zeroPointsT", |
| "semantic": "zero_points", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$zeroPointElement" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputElement" |
| } |
| ], |
| "q4NoZero": [ |
| { |
| "name": "data", |
| "arg": "dataT", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$dataElement" |
| }, |
| { |
| "name": "indices", |
| "arg": "indicesT", |
| "semantic": "indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$indexScalar" |
| }, |
| { |
| "name": "scales", |
| "arg": "scalesT", |
| "semantic": "scales", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scaleScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputElement" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" }, |
| { "name": "packedCols", "type": "u32", "value": "dim(shapes.data, 1)" }, |
| { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" }, |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" } |
| ] |
| } |
| } |
| ], |
| "q4Zero": [ |
| { |
| "name": "data", |
| "arg": "dataT", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$dataElement" |
| }, |
| { |
| "name": "indices", |
| "arg": "indicesT", |
| "semantic": "indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$indexScalar" |
| }, |
| { |
| "name": "scales", |
| "arg": "scalesT", |
| "semantic": "scales", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scaleScalar" |
| }, |
| { |
| "name": "zero_points", |
| "arg": "zeroPointsT", |
| "semantic": "zero_points", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$zeroPointElement" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputElement" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" }, |
| { "name": "packedCols", "type": "u32", "value": "dim(shapes.data, 1)" }, |
| { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" }, |
| { "name": "zeroPointCols", "type": "u32", "value": "zeroPointCols" }, |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" } |
| ] |
| } |
| } |
| ], |
| "q8Vec4NoZero": [ |
| { |
| "name": "data", |
| "arg": "dataT", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$dataElement" |
| }, |
| { |
| "name": "indices", |
| "arg": "indicesT", |
| "semantic": "indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$indexScalar" |
| }, |
| { |
| "name": "scales", |
| "arg": "scalesT", |
| "semantic": "scales", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scaleScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputElement" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" }, |
| { "name": "cols", "type": "u32", "value": "dim(shapes.output, 1)" }, |
| { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" }, |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" } |
| ] |
| } |
| } |
| ], |
| "q8Vec4Zero": [ |
| { |
| "name": "data", |
| "arg": "dataT", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$dataElement" |
| }, |
| { |
| "name": "indices", |
| "arg": "indicesT", |
| "semantic": "indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$indexScalar" |
| }, |
| { |
| "name": "scales", |
| "arg": "scalesT", |
| "semantic": "scales", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scaleScalar" |
| }, |
| { |
| "name": "zero_points", |
| "arg": "zeroPointsT", |
| "semantic": "zero_points", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$zeroPointElement" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputElement" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" }, |
| { "name": "cols", "type": "u32", "value": "dim(shapes.output, 1)" }, |
| { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" }, |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" } |
| ] |
| } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "q8_no_zero_vec4", |
| "priority": 10, |
| "when": ["q8ShapeValid", "noZeroMode", "bits == 8", "blockSize % 4 == 0", "dim(shapes.output, 1) % 4 == 0", "workgroupFits", "foldedDispatchFits"], |
| "constants": { |
| "hasZero": false, |
| "scalarTail": false, |
| "workgroupSize": "workgroupSize", |
| "dataElement": "\"vec4<u32>\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"vec4<f32>\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q8-vec4.wgsl.jinja", |
| "bindings": "q8Vec4NoZero", |
| "dispatch": { |
| "threads": "dim(shapes.indices, 0) * (dim(shapes.output, 1) / 4)", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "q4_no_zero_pair", |
| "priority": 10, |
| "when": ["q4ShapeValid", "noZeroMode", "bits == 4", "workgroupFits", "foldedDispatchFits"], |
| "constants": { |
| "hasZero": false, |
| "workgroupSize": "workgroupSize", |
| "dataElement": "\"u32\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"vec2<f32>\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q4-pair.wgsl.jinja", |
| "bindings": "q4NoZero", |
| "dispatch": { |
| "threads": "dim(shapes.indices, 0) * dim(shapes.data, 1)", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "q4_zero_pair", |
| "priority": 10, |
| "when": ["q4ShapeValid", "zeroMode", "bits == 4", "workgroupFits", "foldedDispatchFits"], |
| "constants": { |
| "hasZero": true, |
| "workgroupSize": "workgroupSize", |
| "dataElement": "\"u32\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"vec2<f32>\"", |
| "zeroPointElement": "\"u32\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q4-pair.wgsl.jinja", |
| "bindings": "q4Zero", |
| "dispatch": { |
| "threads": "dim(shapes.indices, 0) * dim(shapes.data, 1)", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "q8_zero_vec4", |
| "priority": 10, |
| "when": ["q8ShapeValid", "zeroMode", "bits == 8", "blockSize % 4 == 0", "dim(shapes.output, 1) % 4 == 0", "workgroupFits", "foldedDispatchFits"], |
| "constants": { |
| "hasZero": true, |
| "scalarTail": false, |
| "workgroupSize": "workgroupSize", |
| "dataElement": "\"vec4<u32>\"", |
| "zeroPointElement": "\"u32\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"vec4<f32>\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q8-vec4.wgsl.jinja", |
| "bindings": "q8Vec4Zero", |
| "dispatch": { |
| "threads": "dim(shapes.indices, 0) * (dim(shapes.output, 1) / 4)", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "q8_no_zero_tail4", |
| "priority": 5, |
| "when": ["q8ShapeValid", "noZeroMode", "bits == 8", "workgroupFits", "foldedDispatchFits"], |
| "constants": { |
| "hasZero": false, |
| "scalarTail": true, |
| "workgroupSize": "workgroupSize", |
| "dataElement": "\"u32\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"f32\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q8-vec4.wgsl.jinja", |
| "bindings": "noZero", |
| "dispatch": { |
| "threads": "dim(shapes.indices, 0) * ceilDiv(dim(shapes.output, 1), 4)", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "q8_zero_tail4", |
| "priority": 5, |
| "when": ["q8ShapeValid", "zeroMode", "bits == 8", "workgroupFits", "foldedDispatchFits"], |
| "constants": { |
| "hasZero": true, |
| "scalarTail": true, |
| "workgroupSize": "workgroupSize", |
| "dataElement": "\"u32\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"f32\"", |
| "zeroPointElement": "\"u32\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q8-vec4.wgsl.jinja", |
| "bindings": "zero", |
| "dispatch": { |
| "threads": "dim(shapes.indices, 0) * ceilDiv(dim(shapes.output, 1), 4)", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
