{ "domain": "com.microsoft", "name": "MatMulNBits", "sinceVersion": 1, "description": "Matrix multiplication with `B` block-quantized along K and dequantized as `(code - zero_point) * scale`. Each power-of-two `block_size` group has a scale and optional zero point; optional bias is added afterward. Two-, four-, and eight-bit codes are packed low-first, and `A` may have rank 2 or 3. This package supports standard unpacked zero points with the same dtype as `A`. Deprecated `g_idx`, prepacked weights, and bfloat16 tensors are not implemented.", "inputs": [ { "role": "A", "dtype": "T1", "description": "Float input matrix, not quantized. Rank 2 has shape `(M, K)` and rank 3 has shape `(batch, sequence, K)`; only the last axis is the reduction axis and the leading axes fold into the row count, so the ordinary activation needs no surrounding Reshape." }, { "role": "B", "dtype": "uint8", "rank": 3, "description": "Bit-packed uint8 weight matrix of shape `(N, k_blocks, blob_size)`, where `k_blocks = ceil(K / block_size)` and `blob_size = block_size * bits / 8`. Codes are packed low-first along K." }, { "role": "scales", "dtype": "T1", "rank": 2, "description": "Per-block dequantization scale factors of shape `(N, k_blocks)`, with the same dtype as `A`." }, { "role": "zero_points", "dtype": "T3", "rank": 2, "optional": true, "description": "Standard unpacked per-block zero points with shape `(N, k_blocks)` and the same dtype as `A`. Omission uses `2^(bits - 1)`." }, { "role": "bias", "dtype": "T1", "rank": 1, "optional": true, "description": "Optional bias vector of shape `[N]` added to the output." } ], "outputs": [ { "role": "Y", "dtype": "T1", "rank": "ranks.A", "shape": "shapes.A[:-1] + [attrs.N]", "description": "Result of A multiplied by the dequantized weight matrix, with optional bias, same dtype and rank as A: the leading axes of A with a trailing N." } ], "attributes": { "accuracy_level": 0, "bits": 4 }, "attributeDescriptions": { "K": "Input feature dimension of the weight matrix.", "N": "Output feature dimension of the weight matrix.", "accuracy_level": "Minimum internal accuracy level: 0 (unset), 1 (float32), 2 (float16), 3 (bfloat16), or 4 (int8).", "bits": "Bit width used to quantize B; this package supports 2, 4, and 8.", "block_size": "Power-of-two quantization block size along K; it must be at least 16." }, "attributeConstraints": { "K": { "required": true }, "N": { "required": true }, "accuracy_level": { "values": [0, 1, 2, 3, 4] }, "bits": { "values": [2, 4, 8] }, "block_size": { "required": true } }, "typeConstraints": { "T1": ["float32", "float16"], "T3": ["float32", "float16"] }, "args": { "aT": { "kind": "tensor", "semantic": "A", "role": "input" }, "bT": { "kind": "tensor", "semantic": "B", "role": "input" }, "scalesT": { "kind": "tensor", "semantic": "scales", "role": "input" }, "zeroPointsT": { "kind": "tensor", "semantic": "zero_points", "role": "input", "required": false }, "biasT": { "kind": "tensor", "semantic": "bias", "role": "input", "required": false }, "yT": { "kind": "tensor", "semantic": "Y", "role": "output" } }, "tunables": { "REGISTER_TILE_TAILK_MIN_N": 256, "WORKGROUP_SIZE": 64, "GEMV_N_COLS": 4, "REGISTER_TILE_MIN_N": 1024, "REGISTER_TILE_TALL_MIN_M": 512, "REGISTER_TILE_TALL_MIN_N": 256, "REGISTER_TILE_MIN_WORKGROUPS": 64, "REGISTER_TILE_LARGE_M": 64, "REGISTER_TILE_BK32_MIN_M": 128, "REGISTER_TILE_MIN_M": 16 }, "constants": { "B_LEN": "attrs.N * kBlocksExpected * blobSizeExpected", "SCALES_LEN": "attrs.N * kBlocksExpected", "BIAS_LEN": "attrs.N" }, "bindingSets": { "dp4aQuantize": [ { "name": "a", "arg": "aT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "a_quant", "semantic": "aQuant", "buffer": { "type": "storage" }, "elementType": "u32" }, { "name": "a_scales", "semantic": "aScales", "buffer": { "type": "storage" }, "elementType": "f32" } ], "dp4aGemm": [ { "name": "a_quant", "semantic": "aQuant", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, { "name": "a_scales", "semantic": "aScales", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "b", "arg": "bT", "buffer": { "type": "read-only-storage" }, "elementType": "u32", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "buffer": { "type": "read-only-storage" }, "elementType": "f32", "length": "$SCALES_LEN" }, { "name": "y", "arg": "yT", "buffer": { "type": "storage" }, "elementType": "f32" } ], "genericZeroBias": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "zero_points", "arg": "zeroPointsT", "semantic": "zero_points", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$SCALES_LEN" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$BIAS_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "aRows" }, { "name": "K", "type": "u32", "value": "attrs.K" }, { "name": "N", "type": "u32", "value": "attrs.N" }, { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }, { "name": "kBlocks", "type": "u32", "value": "dim(shapes.B, 1)" }, { "name": "blobSize", "type": "u32", "value": "dim(shapes.B, 2)" } ] } } ], "genericZeroOnly": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "zero_points", "arg": "zeroPointsT", "semantic": "zero_points", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$SCALES_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "aRows" }, { "name": "K", "type": "u32", "value": "attrs.K" }, { "name": "N", "type": "u32", "value": "attrs.N" }, { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }, { "name": "kBlocks", "type": "u32", "value": "dim(shapes.B, 1)" }, { "name": "blobSize", "type": "u32", "value": "dim(shapes.B, 2)" } ] } } ], "genericBiasOnly": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$BIAS_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "aRows" }, { "name": "K", "type": "u32", "value": "attrs.K" }, { "name": "N", "type": "u32", "value": "attrs.N" }, { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }, { "name": "kBlocks", "type": "u32", "value": "dim(shapes.B, 1)" }, { "name": "blobSize", "type": "u32", "value": "dim(shapes.B, 2)" } ] } } ], "genericDefaultZero": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "aRows" }, { "name": "K", "type": "u32", "value": "attrs.K" }, { "name": "N", "type": "u32", "value": "attrs.N" }, { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }, { "name": "kBlocks", "type": "u32", "value": "dim(shapes.B, 1)" }, { "name": "blobSize", "type": "u32", "value": "dim(shapes.B, 2)" } ] } } ], "gemvZeroBias": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "zero_points", "arg": "zeroPointsT", "semantic": "zero_points", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$SCALES_LEN" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$BIAS_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" }, { "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" }, { "name": "kBlocks", "type": "u32", "value": "dim(shapes.B, 1)" }, { "name": "blobSize", "type": "u32", "value": "dim(shapes.B, 2)" } ] } } ], "gemvZeroOnly": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "zero_points", "arg": "zeroPointsT", "semantic": "zero_points", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$SCALES_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" }, { "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" }, { "name": "kBlocks", "type": "u32", "value": "dim(shapes.B, 1)" }, { "name": "blobSize", "type": "u32", "value": "dim(shapes.B, 2)" } ] } } ], "gemvBiasOnly": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$BIAS_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" }, { "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" }, { "name": "kBlocks", "type": "u32", "value": "dim(shapes.B, 1)" }, { "name": "blobSize", "type": "u32", "value": "dim(shapes.B, 2)" } ] } } ], "gemvDefaultZero": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" }, { "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" }, { "name": "kBlocks", "type": "u32", "value": "dim(shapes.B, 1)" }, { "name": "blobSize", "type": "u32", "value": "dim(shapes.B, 2)" } ] } } ], "prefillZeroBias": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "zero_points", "arg": "zeroPointsT", "semantic": "zero_points", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$SCALES_LEN" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$BIAS_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" } ], "prefillZeroOnly": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "zero_points", "arg": "zeroPointsT", "semantic": "zero_points", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$SCALES_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" } ], "prefillBiasOnly": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$aScalar", "length": "$BIAS_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" } ], "prefillDefaultZero": [ { "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": "$bScalar", "length": "$B_LEN" }, { "name": "scales", "arg": "scalesT", "semantic": "scales", "buffer": { "type": "read-only-storage" }, "elementType": "$scaleScalar", "length": "$SCALES_LEN" }, { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputScalar" } ] }, "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", "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16", "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", "packedFeature": "device.wgslLanguageFeatures.has(\"packed_4x8_integer_dot_product\")", "kBlocksExpected": "ceilDiv(attrs.K, attrs.block_size)", "blobSizeExpected": "ceilDiv(attrs.block_size * attrs.bits, 8)", "aRows": "numel(shapes.A) / max(1, attrs.K)", "aRankOk": "(ranks.A == 2 or ranks.A == 3) and ranks.Y == ranks.A and dim(shapes.A, ranks.A - 1) == attrs.K and dim(shapes.Y, ranks.Y - 1) == attrs.N and dim(shapes.Y, 0) == dim(shapes.A, 0) and (ranks.A == 2 or dim(shapes.Y, 1) == dim(shapes.A, 1))", "dispatchN4": "ceilDiv(attrs.N, 4)", "gemvDispatchN": "ceilDiv(attrs.N, tunables.GEMV_N_COLS)", "dispatchN32": "ceilDiv(attrs.N, 32)", "dispatchN64": "ceilDiv(attrs.N, 64)", "dispatchM32": "ceilDiv(aRows, 32)", "dispatchM64": "ceilDiv(aRows, 64)", "sgmatTileRows": "64 if aRows >= 64 else 32", "sgmatWorkgroupSize": "256 if aRows >= 64 else 128", "sgmatRowSubtiles": "4 if aRows >= 64 else 2", "sgmatBLoadsPerRow": "sgmatRowSubtiles", "sgmatBLoadWidth": "8 if aRows >= 64 else 16", "sgmatNumSubgroups": "8 if aRows >= 64 else 4", "sgmatWorkgroupStorageBytes": "24576 if aRows >= 64 else 16384", "sgmatDispatchM": "dispatchM64 if aRows >= 64 else dispatchM32", "tiledRegBK": "32 if aRows >= tunables.REGISTER_TILE_BK32_MIN_M else 16", "aFloatOk": "(tensorDtypes.A == \"float32\" or tensorDtypes.A == \"float16\") and f16Ok(tensorDtypes.A)", "portableWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", "bitsSupported": "attrs.bits == 2 or attrs.bits == 4 or attrs.bits == 8", "blockSizeSupported": "attrs.block_size >= 16 and attrs.block_size == pow2ceil(attrs.block_size)", "commonShapeValid": "aRankOk and ranks.B == 3 and ranks.scales == 2 and aFloatOk and blockSizeSupported and tensorDtypes.B == \"uint8\" and tensorDtypes.scales == tensorDtypes.A and tensorDtypes.Y == tensorDtypes.A and attrs.K > 0 and attrs.N > 0 and dim(shapes.B, 0) == attrs.N and dim(shapes.B, 1) == kBlocksExpected and dim(shapes.B, 2) == blobSizeExpected and dim(shapes.scales, 0) == attrs.N and dim(shapes.scales, 1) == dim(shapes.B, 1)", "gemvShapeValid": "commonShapeValid and aRows == 1", "zeroPointsValid": "present.zeroPointsT and ranks.zero_points == 2 and tensorDtypes.zero_points == tensorDtypes.A and dim(shapes.zero_points, 0) == attrs.N and dim(shapes.zero_points, 1) == dim(shapes.B, 1)", "biasValid": "present.biasT and ranks.bias == 1 and tensorDtypes.bias == tensorDtypes.A and dim(shapes.bias, 0) == attrs.N", "defaultEpilogue": "not present.zeroPointsT and not present.biasT", "zeroBiasEpilogue": "zeroPointsValid and biasValid", "zeroOnlyEpilogue": "zeroPointsValid and not present.biasT", "biasOnlyEpilogue": "not present.zeroPointsT and biasValid", "portableWorkgroupFits": "portableWorkgroupSize > 0 and portableWorkgroupSize * 64 <= device.limits.maxComputeWorkgroupStorageSize", "tiledWorkgroupFits": "16 <= device.limits.maxComputeWorkgroupSizeX and 16 <= device.limits.maxComputeWorkgroupSizeY and 256 <= device.limits.maxComputeInvocationsPerWorkgroup and 4096 <= device.limits.maxComputeWorkgroupStorageSize", "tiledRegWorkgroupFits": "tiledWorkgroupFits and 16384 <= device.limits.maxComputeWorkgroupStorageSize", "mediumTiledRegWorkgroupFits": "tiledWorkgroupFits and 6144 <= device.limits.maxComputeWorkgroupStorageSize", "sgmatWorkgroupFits": "sgmatWorkgroupSize <= deviceWorkgroupCap and sgmatWorkgroupStorageBytes <= device.limits.maxComputeWorkgroupStorageSize", "registerTileShape": "aRows >= tunables.REGISTER_TILE_MIN_M and ((attrs.N >= tunables.REGISTER_TILE_MIN_N or (aRows >= tunables.REGISTER_TILE_TALL_MIN_M and attrs.N >= tunables.REGISTER_TILE_TALL_MIN_N) or (attrs.K % attrs.block_size != 0 and attrs.N >= tunables.REGISTER_TILE_TAILK_MIN_N)) and (attrs.K % attrs.block_size != 0 or aRows >= tunables.REGISTER_TILE_LARGE_M or dispatchM64 * dispatchN64 >= tunables.REGISTER_TILE_MIN_WORKGROUPS))", "portableTile4Preferred": "registerTileShape and attrs.K % attrs.block_size == 0 and (not device.features.has(\"subgroups\") or narrowSubgroupRange) and has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize * 2 < portableWorkgroupSize", "portableMediumRegisterPreferred": "portableTile4Preferred and aRows >= 128 and attrs.K >= 128", "registerTilePreferred": "registerTileShape and not portableTile4Preferred", "mediumRegisterEligible": "registerTilePreferred or portableMediumRegisterPreferred", "largeTiledRegEligible": "registerTilePreferred and tiledRegWorkgroupFits and dispatchM64 <= device.limits.maxComputeWorkgroupsPerDimension", "mediumTiledRegEligible": "mediumRegisterEligible and mediumTiledRegWorkgroupFits and dispatchM32 <= device.limits.maxComputeWorkgroupsPerDimension", "tiledRegVariantEligible": "largeTiledRegEligible or mediumTiledRegEligible", "tiledRegSelectedBK": "tiledRegBK if largeTiledRegEligible else 16", "tiledRegSelectedTileRows": "64 if largeTiledRegEligible else 32", "tiledRegSelectedThreadRows": "4 if largeTiledRegEligible else 2", "tiledRegSelectedDispatchM": "dispatchM64 if largeTiledRegEligible else dispatchM32" }, "variants": [ { "id": "q4_dp4a_prefill", "priority": 19, "when": ["packedFeature", "commonShapeValid", "defaultEpilogue", "attrs.bits == 4", "attrs.accuracy_level == 4", "tensorDtypes.A == \"float32\"", "attrs.block_size % 32 == 0", "attrs.K % 128 == 0", "attrs.N % 16 == 0", "aRows >= 32", "ceilDiv(attrs.N, 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(aRows, 64) <= device.limits.maxComputeWorkgroupsPerDimension", "16 <= device.limits.maxComputeWorkgroupSizeX", "16 <= device.limits.maxComputeWorkgroupSizeY", "64 <= device.limits.maxComputeWorkgroupSizeX", "256 <= device.limits.maxComputeInvocationsPerWorkgroup", "4608 <= device.limits.maxComputeWorkgroupStorageSize"], "demoteWhen": ["has(device.adapterInfo, \"architecture\") and device.adapterInfo.architecture == \"maxwell\"", "device.features.has(\"chromium-experimental-subgroup-matrix\")", "device.adapterInfo.vendor == \"apple\""], "constants": { "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "vec4Count": "aRows * attrs.K / 4", "blockCount": "aRows * attrs.K / 128" }, "intermediates": [ { "id": "aQuant", "dtype": "uint32", "shape": "[aRows * attrs.K / 4]" }, { "id": "aScales", "dtype": "float32", "shape": "[aRows * attrs.K / 128]" } ], "passes": [ { "id": "quantize", "name": "MatMulNBits.Dp4aQuantizeA", "shader": "matmul-nbits-dp4a-quantize.wgsl.jinja", "bindings": "dp4aQuantize", "dispatch": { "threads": "aRows * attrs.K / 4", "workgroupSize": 64 } }, { "id": "main", "name": "MatMulNBits.Dp4aPrefill", "shader": "matmul-nbits-q4-dp4a-prefill.wgsl.jinja", "bindings": "dp4aGemm", "dispatch": { "x": "ceilDiv(attrs.N, 64)", "y": "ceilDiv(aRows, 64)" } } ] }, { "id": "gemv_default_zero", "priority": 20, "when": ["gemvShapeValid", "defaultEpilogue", "bitsSupported", "portableWorkgroupFits"], "constants": { "gemvNCols": "tunables.GEMV_N_COLS", "hasZero": false, "hasBias": false, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-gemv-q4.wgsl.jinja", "bindings": "gemvDefaultZero", "dispatch": { "workgroups": "gemvDispatchN" } } ] }, { "id": "prefill_sgmat_default_zero", "priority": 18, "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "limits": { "maxComputeWorkgroupStorageSize": 16384 }, "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "wave32Effective", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"], "constants": { "hasZero": false, "hasBias": false, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"", "tileRows": "sgmatTileRows", "workgroupSize": "sgmatWorkgroupSize", "rowSubtiles": "sgmatRowSubtiles", "bLoadsPerRow": "sgmatBLoadsPerRow", "bLoadWidth": "sgmatBLoadWidth", "numSubgroups": "sgmatNumSubgroups" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-sgmat.wgsl.jinja", "bindings": "prefillDefaultZero", "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" } } ] }, { "id": "prefill_tiled_reg_default_zero", "priority": 17, "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "tiledRegVariantEligible"], "constants": { "hasZero": false, "hasBias": false, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"", "bk": "tiledRegSelectedBK", "tileRows": "tiledRegSelectedTileRows", "tileCols": 64, "threadRows": "tiledRegSelectedThreadRows", "threadCols": 4, "alignedBlockLoads": "aRows >= 128 and attrs.block_size % tiledRegSelectedBK == 0" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja", "bindings": "prefillDefaultZero", "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" } } ] }, { "id": "prefill_tiled_default_zero", "priority": 16, "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= device.limits.maxComputeWorkgroupsPerDimension", "dispatchM32 <= device.limits.maxComputeWorkgroupsPerDimension", "tiledWorkgroupFits"], "constants": { "hasZero": false, "hasBias": false, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja", "bindings": "prefillDefaultZero", "dispatch": { "x": "dispatchN32", "y": "dispatchM32" } } ] }, { "id": "prefill_tile4x4_default_zero", "priority": 15, "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"], "constants": { "hasZero": false, "hasBias": false, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja", "bindings": "genericDefaultZero", "dispatch": { "workgroups": "dispatchN4", "y": "min(ceilDiv(aRows, 4), device.limits.maxComputeWorkgroupsPerDimension)" } } ] }, { "id": "default_zero", "priority": 0, "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "portableWorkgroupFits"], "constants": { "hasZero": false, "hasBias": false, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits.wgsl.jinja", "bindings": "genericDefaultZero", "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "constants.workgroupSize" } } ] }, { "id": "gemv_zero_bias", "priority": 20, "when": ["gemvShapeValid", "zeroBiasEpilogue", "bitsSupported", "portableWorkgroupFits"], "constants": { "gemvNCols": "tunables.GEMV_N_COLS", "hasZero": true, "hasBias": true, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-gemv-q4.wgsl.jinja", "bindings": "gemvZeroBias", "dispatch": { "workgroups": "gemvDispatchN" } } ] }, { "id": "prefill_sgmat_zero_bias", "priority": 18, "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "limits": { "maxComputeWorkgroupStorageSize": 16384 }, "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "wave32Effective", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"], "constants": { "hasZero": true, "hasBias": true, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"", "tileRows": "sgmatTileRows", "workgroupSize": "sgmatWorkgroupSize", "rowSubtiles": "sgmatRowSubtiles", "bLoadsPerRow": "sgmatBLoadsPerRow", "bLoadWidth": "sgmatBLoadWidth", "numSubgroups": "sgmatNumSubgroups" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-sgmat.wgsl.jinja", "bindings": "prefillZeroBias", "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" } } ] }, { "id": "prefill_tiled_reg_zero_bias", "priority": 17, "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "tiledRegVariantEligible"], "constants": { "hasZero": true, "hasBias": true, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"", "bk": "tiledRegSelectedBK", "tileRows": "tiledRegSelectedTileRows", "tileCols": 64, "threadRows": "tiledRegSelectedThreadRows", "threadCols": 4, "alignedBlockLoads": "aRows >= 128 and attrs.block_size % tiledRegSelectedBK == 0" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja", "bindings": "prefillZeroBias", "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" } } ] }, { "id": "prefill_tiled_zero_bias", "priority": 16, "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= device.limits.maxComputeWorkgroupsPerDimension", "dispatchM32 <= device.limits.maxComputeWorkgroupsPerDimension", "tiledWorkgroupFits"], "constants": { "hasZero": true, "hasBias": true, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja", "bindings": "prefillZeroBias", "dispatch": { "x": "dispatchN32", "y": "dispatchM32" } } ] }, { "id": "prefill_tile4x4_zero_bias", "priority": 15, "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"], "constants": { "hasZero": true, "hasBias": true, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja", "bindings": "genericZeroBias", "dispatch": { "workgroups": "dispatchN4", "y": "min(ceilDiv(aRows, 4), device.limits.maxComputeWorkgroupsPerDimension)" } } ] }, { "id": "zero_bias", "priority": 0, "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "portableWorkgroupFits"], "constants": { "hasZero": true, "hasBias": true, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits.wgsl.jinja", "bindings": "genericZeroBias", "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "constants.workgroupSize" } } ] }, { "id": "gemv_zero_only", "priority": 20, "when": ["gemvShapeValid", "zeroOnlyEpilogue", "bitsSupported", "portableWorkgroupFits"], "constants": { "gemvNCols": "tunables.GEMV_N_COLS", "hasZero": true, "hasBias": false, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-gemv-q4.wgsl.jinja", "bindings": "gemvZeroOnly", "dispatch": { "workgroups": "gemvDispatchN" } } ] }, { "id": "prefill_sgmat_zero_only", "priority": 18, "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "limits": { "maxComputeWorkgroupStorageSize": 16384 }, "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "wave32Effective", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"], "constants": { "hasZero": true, "hasBias": false, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"", "tileRows": "sgmatTileRows", "workgroupSize": "sgmatWorkgroupSize", "rowSubtiles": "sgmatRowSubtiles", "bLoadsPerRow": "sgmatBLoadsPerRow", "bLoadWidth": "sgmatBLoadWidth", "numSubgroups": "sgmatNumSubgroups" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-sgmat.wgsl.jinja", "bindings": "prefillZeroOnly", "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" } } ] }, { "id": "prefill_tiled_reg_zero_only", "priority": 17, "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "tiledRegVariantEligible"], "constants": { "hasZero": true, "hasBias": false, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"", "bk": "tiledRegSelectedBK", "tileRows": "tiledRegSelectedTileRows", "tileCols": 64, "threadRows": "tiledRegSelectedThreadRows", "threadCols": 4, "alignedBlockLoads": "aRows >= 128 and attrs.block_size % tiledRegSelectedBK == 0" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja", "bindings": "prefillZeroOnly", "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" } } ] }, { "id": "prefill_tiled_zero_only", "priority": 16, "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= device.limits.maxComputeWorkgroupsPerDimension", "dispatchM32 <= device.limits.maxComputeWorkgroupsPerDimension", "tiledWorkgroupFits"], "constants": { "hasZero": true, "hasBias": false, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja", "bindings": "prefillZeroOnly", "dispatch": { "x": "dispatchN32", "y": "dispatchM32" } } ] }, { "id": "prefill_tile4x4_zero_only", "priority": 15, "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"], "constants": { "hasZero": true, "hasBias": false, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja", "bindings": "genericZeroOnly", "dispatch": { "workgroups": "dispatchN4", "y": "min(ceilDiv(aRows, 4), device.limits.maxComputeWorkgroupsPerDimension)" } } ] }, { "id": "zero_only", "priority": 0, "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "portableWorkgroupFits"], "constants": { "hasZero": true, "hasBias": false, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits.wgsl.jinja", "bindings": "genericZeroOnly", "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "constants.workgroupSize" } } ] }, { "id": "gemv_bias_only", "priority": 20, "when": ["gemvShapeValid", "biasOnlyEpilogue", "bitsSupported", "portableWorkgroupFits"], "constants": { "gemvNCols": "tunables.GEMV_N_COLS", "hasZero": false, "hasBias": true, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-gemv-q4.wgsl.jinja", "bindings": "gemvBiasOnly", "dispatch": { "workgroups": "gemvDispatchN" } } ] }, { "id": "prefill_sgmat_bias_only", "priority": 18, "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "limits": { "maxComputeWorkgroupStorageSize": 16384 }, "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "wave32Effective", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"], "constants": { "hasZero": false, "hasBias": true, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"", "tileRows": "sgmatTileRows", "workgroupSize": "sgmatWorkgroupSize", "rowSubtiles": "sgmatRowSubtiles", "bLoadsPerRow": "sgmatBLoadsPerRow", "bLoadWidth": "sgmatBLoadWidth", "numSubgroups": "sgmatNumSubgroups" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-sgmat.wgsl.jinja", "bindings": "prefillBiasOnly", "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" } } ] }, { "id": "prefill_tiled_reg_bias_only", "priority": 17, "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "tiledRegVariantEligible"], "constants": { "hasZero": false, "hasBias": true, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"", "bk": "tiledRegSelectedBK", "tileRows": "tiledRegSelectedTileRows", "tileCols": 64, "threadRows": "tiledRegSelectedThreadRows", "threadCols": 4, "alignedBlockLoads": "aRows >= 128 and attrs.block_size % tiledRegSelectedBK == 0" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja", "bindings": "prefillBiasOnly", "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" } } ] }, { "id": "prefill_tiled_bias_only", "priority": 16, "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= device.limits.maxComputeWorkgroupsPerDimension", "dispatchM32 <= device.limits.maxComputeWorkgroupsPerDimension", "tiledWorkgroupFits"], "constants": { "hasZero": false, "hasBias": true, "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "M": "aRows", "K": "attrs.K", "N": "attrs.N", "kBlocks": "dim(shapes.B, 1)", "blockSize": "attrs.block_size", "blobSize": "dim(shapes.B, 2)", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja", "bindings": "prefillBiasOnly", "dispatch": { "x": "dispatchN32", "y": "dispatchM32" } } ] }, { "id": "prefill_tile4x4_bias_only", "priority": 15, "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"], "constants": { "hasZero": false, "hasBias": true, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja", "bindings": "genericBiasOnly", "dispatch": { "workgroups": "dispatchN4", "y": "min(ceilDiv(aRows, 4), device.limits.maxComputeWorkgroupsPerDimension)" } } ] }, { "id": "bias_only", "priority": 0, "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "portableWorkgroupFits"], "constants": { "hasZero": false, "hasBias": true, "workgroupSize": "portableWorkgroupSize", "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bScalar": "\"u32\"", "scaleScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"", "bits": "attrs.bits", "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")", "usesF16": "tensorDtypes.A == \"float16\"" }, "passes": [ { "id": "main", "shader": "matmul-nbits.wgsl.jinja", "bindings": "genericBiasOnly", "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "constants.workgroupSize" } } ] } ] }