| { |
| "domain": "com.microsoft", |
| "name": "MatMulNBitsQkv", |
| "sinceVersion": 1, |
| "inputs": { |
| "aT": { "onnx": "A", "dtype": "T1" }, |
| "skipT": { "onnx": "skip", "dtype": "T1", "optional": true }, |
| "normScaleT": { "onnx": "norm_scale", "dtype": "T1", "rank": 1 }, |
| "qBT": { "onnx": "q_B", "dtype": "T2", "rank": 3, "layout": "packed" }, |
| "qScalesT": { "onnx": "q_scales", "dtype": "T1", "rank": 2 }, |
| "kBT": { "onnx": "k_B", "dtype": "T2", "rank": 3, "layout": "packed" }, |
| "kScalesT": { "onnx": "k_scales", "dtype": "T1", "rank": 2 }, |
| "vBT": { "onnx": "v_B", "dtype": "T2", "rank": 3, "layout": "packed" }, |
| "vScalesT": { "onnx": "v_scales", "dtype": "T1", "rank": 2 } |
| }, |
| "outputs": { |
| "qT": { "onnx": "Q", "dtype": "T1", "rank": "ranks.aT", "shape": "shapes.aT[:-1] + [attrs.Nq]" }, |
| "kT": { "onnx": "K", "dtype": "T1", "rank": "ranks.aT", "shape": "shapes.aT[:-1] + [attrs.Nkv]" }, |
| "vT": { "onnx": "V", "dtype": "T1", "rank": "ranks.aT", "shape": "shapes.aT[:-1] + [attrs.Nkv]" }, |
| "residualT": { |
| "onnx": "input_skip_bias_sum", |
| "dtype": "T1", |
| "rank": "ranks.aT", |
| "optional": true, |
| "shape": "shapes.aT" |
| } |
| }, |
| "attributes": { |
| "accuracy_level": { "default": 0 }, |
| "bits": { "default": 4 }, |
| "epsilon": { "default": 9.999999974752427e-7 }, |
| "K": {}, |
| "Nq": {}, |
| "Nkv": {}, |
| "block_size": {} |
| }, |
| "attributeConstraints": { |
| "K": { "required": true }, |
| "Nq": { "required": true }, |
| "Nkv": { "required": true }, |
| "accuracy_level": { "values": [0] }, |
| "bits": { "values": [4] }, |
| "block_size": { "required": true, "values": [32] } |
| }, |
| "typeConstraints": { "T1": ["float32", "float16"], "T2": ["uint8"] }, |
| "tunables": { |
| "TILE_N": { "default": 8 }, |
| "LANES": { "default": 8 }, |
| "NORM_WORKGROUP_SIZE": { "default": 128 }, |
| "ROW_TILE": { "default": 8 }, |
| "DECODE_WORKGROUP_SIZE": { "default": 64 } |
| }, |
| "derive": { |
| "aRows": "numel(shapes.aT) / max(1, attrs.K)", |
| "rowTile": "1 if aRows <= 1 else min(aRows, tunables.ROW_TILE)", |
| "rowGroups": "ceilDiv(aRows, rowTile)", |
| "kBlocks": "dim(shapes.qBT, 1)", |
| "blobSize": "dim(shapes.qBT, 2)", |
| "codesPerByte": "8 / attrs.bits", |
| "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)", |
| "pairSharesWord": "codesPerByte >= 2", |
| "epsilonValue": "attrs.epsilon", |
| "weightShapeOk": "ranks.qBT == 3 and ranks.kBT == 3 and ranks.vBT == 3 and dim(shapes.qBT, 0) == attrs.Nq and dim(shapes.kBT, 0) == attrs.Nkv and dim(shapes.vBT, 0) == attrs.Nkv and dim(shapes.kBT, 1) == kBlocks and dim(shapes.vBT, 1) == kBlocks and dim(shapes.kBT, 2) == blobSize and dim(shapes.vBT, 2) == blobSize and kBlocks == ceilDiv(attrs.K, attrs.block_size) and blobSize * 8 == attrs.block_size * attrs.bits", |
| "scaleShapeOk": "ranks.qScalesT == 2 and ranks.kScalesT == 2 and ranks.vScalesT == 2 and dim(shapes.qScalesT, 0) == attrs.Nq and dim(shapes.qScalesT, 1) == kBlocks and dim(shapes.kScalesT, 0) == attrs.Nkv and dim(shapes.kScalesT, 1) == kBlocks and dim(shapes.vScalesT, 0) == attrs.Nkv and dim(shapes.vScalesT, 1) == kBlocks", |
| "ioShapeOk": "(ranks.aT == 2 or ranks.aT == 3) and dim(shapes.aT, ranks.aT - 1) == attrs.K and ranks.qT == ranks.aT and ranks.kT == ranks.aT and ranks.vT == ranks.aT and dim(shapes.qT, ranks.qT - 1) == attrs.Nq and dim(shapes.kT, ranks.kT - 1) == attrs.Nkv and dim(shapes.vT, ranks.vT - 1) == attrs.Nkv and sameShape(prefix(shapes.qT, ranks.qT - 1), prefix(shapes.aT, ranks.aT - 1)) and sameShape(prefix(shapes.kT, ranks.kT - 1), prefix(shapes.aT, ranks.aT - 1)) and sameShape(prefix(shapes.vT, ranks.vT - 1), prefix(shapes.aT, ranks.aT - 1))", |
| "dtypeOk": "tensorDtypes.qBT == \"uint8\" and tensorDtypes.kBT == \"uint8\" and tensorDtypes.vBT == \"uint8\" and tensorDtypes.qScalesT == tensorDtypes.aT and tensorDtypes.kScalesT == tensorDtypes.aT and tensorDtypes.vScalesT == tensorDtypes.aT and tensorDtypes.qT == tensorDtypes.aT and tensorDtypes.kT == tensorDtypes.aT and tensorDtypes.vT == tensorDtypes.aT and tensorDtypes.normScaleT == tensorDtypes.aT and f16Ok(tensorDtypes.aT)", |
| "lanesPow2": "tunables.LANES == pow2ceil(tunables.LANES)", |
| "normContractOk": "ranks.normScaleT == 1 and dim(shapes.normScaleT, 0) == attrs.K and (sameShape(shapes.skipT, shapes.aT) and tensorDtypes.skipT == tensorDtypes.aT if present.skipT else true) and (sameShape(shapes.residualT, shapes.aT) and tensorDtypes.residualT == tensorDtypes.aT and present.skipT if present.residualT else true)", |
| "qkvShapeOk": "weightShapeOk and scaleShapeOk and ioShapeOk and dtypeOk and lanesPow2 and normContractOk and pairSharesWord and attrs.K > 0 and attrs.Nq > 0 and attrs.Nkv > 0", |
| "decodeWalk": "aRows <= 1", |
| "decodeCols": "4", |
| "gemvWalk": "rowTile == 1 and blobSize % 16 == 0", |
| "decodeActVec4": "gemvWalk and attrs.K % attrs.block_size == 0", |
| "tileCols": "decodeCols if decodeWalk else tunables.TILE_N", |
| "decodeWorkgroupOk": "tunables.DECODE_WORKGROUP_SIZE >= 4 and pow2ceil(tunables.DECODE_WORKGROUP_SIZE) == tunables.DECODE_WORKGROUP_SIZE and tunables.DECODE_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.DECODE_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX", |
| "projectionTiles": "ceilDiv(attrs.Nq, tileCols) + 2 * ceilDiv(attrs.Nkv, tileCols)", |
| "dispatchFits": "decodeWorkgroupOk and projectionTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and aRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and tunables.TILE_N * tunables.LANES <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.TILE_N * tunables.LANES <= device.limits.maxComputeWorkgroupSizeX and tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX", |
| "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"", |
| "scalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"", |
| "K": "attrs.K", |
| "nq": "attrs.Nq", |
| "nkv": "attrs.Nkv", |
| "blockSize": "attrs.block_size", |
| "bits": "attrs.bits", |
| "defaultZero": "\"8.0\"", |
| "tileN": "tunables.TILE_N", |
| "lanes": "tunables.LANES", |
| "hidden": "attrs.K", |
| "workgroupSize": "tunables.NORM_WORKGROUP_SIZE", |
| "epsilon": "epsilonValue", |
| "hasSkip": "present.skipT", |
| "writeResidual": "present.residualT", |
| "K_LEN": "attrs.K", |
| "rowCount": "aRows", |
| "decodeNCols": "decodeCols", |
| "actVec4": "decodeActVec4", |
| "weightElement": "\"vec4<u32>\" if gemvWalk else \"u32\"", |
| "normedElement": "\"vec4<f32>\" if decodeActVec4 else \"f32\"", |
| "decodeWorkgroupSize": "tunables.DECODE_WORKGROUP_SIZE", |
| "useSubgroups": "device.features.has(\"subgroups\")" |
| }, |
| "when": ["dispatchFits", "qkvShapeOk"], |
| "bindings": { |
| "a": { "arg": "aT", "buffer": "read-only-storage", "elementType": "$aScalar" }, |
| "norm_scale": { "arg": "normScaleT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$K_LEN" }, |
| "normed": { "scratch": "normedA", "buffer": "storage", "elementType": "f32" }, |
| "params": { "buffer": "uniform", "struct": [{ "name": "rows", "type": "u32", "value": "aRows" }] }, |
| "skip": { "arg": "skipT", "buffer": "read-only-storage", "elementType": "$aScalar" }, |
| "residual": { "arg": "residualT", "buffer": "storage", "elementType": "$aScalar" }, |
| "normed_2": { |
| "scratch": "normedA", |
| "name": "normed", |
| "buffer": "read-only-storage", |
| "elementType": "$normedElement" |
| }, |
| "q_b": { "arg": "qBT", "buffer": "read-only-storage", "elementType": "$weightElement" }, |
| "q_scales": { "arg": "qScalesT", "buffer": "read-only-storage", "elementType": "$aScalar" }, |
| "k_b": { "arg": "kBT", "buffer": "read-only-storage", "elementType": "$weightElement" }, |
| "k_scales": { "arg": "kScalesT", "buffer": "read-only-storage", "elementType": "$aScalar" }, |
| "v_b": { "arg": "vBT", "buffer": "read-only-storage", "elementType": "$weightElement" }, |
| "v_scales": { "arg": "vScalesT", "buffer": "read-only-storage", "elementType": "$aScalar" }, |
| "q": { "arg": "qT", "buffer": "storage", "elementType": "$aScalar" }, |
| "k": { "arg": "kT", "buffer": "storage", "elementType": "$aScalar" }, |
| "v": { "arg": "vT", "buffer": "storage", "elementType": "$aScalar" } |
| }, |
| "variants": [ |
| { |
| "id": "norm", |
| "priority": 20, |
| "when": ["not present.skipT", "not present.residualT"], |
| "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }], |
| "passes": [ |
| { |
| "id": "norm", |
| "name": "MatMulNBitsQkv.RmsNorm", |
| "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja", |
| "bindings": ["a", "norm_scale", "normed", "params"], |
| "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulNBitsQkv.Projection", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"\"" }, |
| "bindings": ["normed_2", "q_b", "q_scales", "k_b", "k_scales", "v_b", "v_scales", "q", "k", "v"], |
| "dispatch": { "x": "projectionTiles", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "split_norm", |
| "priority": 10, |
| "when": ["not present.skipT", "not present.residualT"], |
| "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }], |
| "passes": [ |
| { |
| "id": "norm", |
| "name": "MatMulNBitsQkv.RmsNorm", |
| "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja", |
| "bindings": ["a", "norm_scale", "normed", "params"], |
| "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 } |
| }, |
| { |
| "id": "q", |
| "name": "MatMulNBitsQkv.ProjectionQ", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"q\"" }, |
| "bindings": ["normed_2", "q_b", "q_scales", "q"], |
| "dispatch": { "x": "ceilDiv(attrs.Nq, tileCols)", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| }, |
| { |
| "id": "k", |
| "name": "MatMulNBitsQkv.ProjectionK", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"k\"" }, |
| "bindings": ["normed_2", "k_b", "k_scales", "k"], |
| "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| }, |
| { |
| "id": "v", |
| "name": "MatMulNBitsQkv.ProjectionV", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"v\"" }, |
| "bindings": ["normed_2", "v_b", "v_scales", "v"], |
| "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "skip", |
| "priority": 20, |
| "when": ["present.skipT", "not present.residualT"], |
| "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }], |
| "passes": [ |
| { |
| "id": "norm", |
| "name": "MatMulNBitsQkv.RmsNorm", |
| "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja", |
| "bindings": ["a", "skip", "norm_scale", "normed", "params"], |
| "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulNBitsQkv.Projection", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"\"" }, |
| "bindings": ["normed_2", "q_b", "q_scales", "k_b", "k_scales", "v_b", "v_scales", "q", "k", "v"], |
| "dispatch": { "x": "projectionTiles", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "split_skip", |
| "priority": 10, |
| "when": ["present.skipT", "not present.residualT"], |
| "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }], |
| "passes": [ |
| { |
| "id": "norm", |
| "name": "MatMulNBitsQkv.RmsNorm", |
| "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja", |
| "bindings": ["a", "skip", "norm_scale", "normed", "params"], |
| "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 } |
| }, |
| { |
| "id": "q", |
| "name": "MatMulNBitsQkv.ProjectionQ", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"q\"" }, |
| "bindings": ["normed_2", "q_b", "q_scales", "q"], |
| "dispatch": { "x": "ceilDiv(attrs.Nq, tileCols)", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| }, |
| { |
| "id": "k", |
| "name": "MatMulNBitsQkv.ProjectionK", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"k\"" }, |
| "bindings": ["normed_2", "k_b", "k_scales", "k"], |
| "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| }, |
| { |
| "id": "v", |
| "name": "MatMulNBitsQkv.ProjectionV", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"v\"" }, |
| "bindings": ["normed_2", "v_b", "v_scales", "v"], |
| "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "skipsum", |
| "priority": 20, |
| "when": ["present.skipT", "present.residualT"], |
| "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }], |
| "passes": [ |
| { |
| "id": "norm", |
| "name": "MatMulNBitsQkv.RmsNorm", |
| "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja", |
| "bindings": ["a", "skip", "norm_scale", "normed", "residual", "params"], |
| "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulNBitsQkv.Projection", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"\"" }, |
| "bindings": ["normed_2", "q_b", "q_scales", "k_b", "k_scales", "v_b", "v_scales", "q", "k", "v"], |
| "dispatch": { "x": "projectionTiles", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "split_skipsum", |
| "priority": 10, |
| "when": ["present.skipT", "present.residualT"], |
| "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }], |
| "passes": [ |
| { |
| "id": "norm", |
| "name": "MatMulNBitsQkv.RmsNorm", |
| "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja", |
| "bindings": ["a", "skip", "norm_scale", "normed", "residual", "params"], |
| "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 } |
| }, |
| { |
| "id": "q", |
| "name": "MatMulNBitsQkv.ProjectionQ", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"q\"" }, |
| "bindings": ["normed_2", "q_b", "q_scales", "q"], |
| "dispatch": { "x": "ceilDiv(attrs.Nq, tileCols)", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| }, |
| { |
| "id": "k", |
| "name": "MatMulNBitsQkv.ProjectionK", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"k\"" }, |
| "bindings": ["normed_2", "k_b", "k_scales", "k"], |
| "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| }, |
| { |
| "id": "v", |
| "name": "MatMulNBitsQkv.ProjectionV", |
| "shader": "qkv-projection.wgsl.jinja", |
| "derive": { "singleProjection": "\"v\"" }, |
| "bindings": ["normed_2", "v_b", "v_scales", "v"], |
| "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| } |
| ] |
| } |
|
|
