{ "domain": "com.microsoft", "name": "QAttention", "sinceVersion": 1, "inputs": { "input": { "dtype": "T1" }, "weight": { "dtype": "T2" }, "bias": { "dtype": "T3" }, "input_scale": { "dtype": "T3" }, "weight_scale": { "dtype": "T3" }, "mask_index": { "dtype": "T4", "optional": true }, "input_zero_point": { "dtype": "T1", "optional": true }, "weight_zero_point": { "dtype": "T2", "optional": true }, "past": { "dtype": "T3", "optional": true } }, "outputs": { "output": { "dtype": "T3", "rank": 3, "shape": ["dim(shapes.input, 0)", "dim(shapes.input, 1)", "dim(shapes.weight, 1) / 3"] }, "present": { "dtype": "T3", "rank": 5, "optional": true, "shape": [2, "dim(shapes.input, 0)", "attrs.num_heads", "(dim(shapes.past, 3) if present.past else 0) + dim(shapes.input, 1)", "dim(shapes.weight, 1) / 3 / attrs.num_heads"] } }, "attributes": { "num_heads": {}, "unidirectional": { "default": 0 }, "do_rotary": { "default": 0 }, "mask_filter_value": { "default": -10000 }, "past_present_share_buffer": {}, "scale": {} }, "attributeConstraints": { "num_heads": { "required": true }, "do_rotary": { "values": [0] }, "past_present_share_buffer": { "values": [0] } }, "typeConstraints": { "T1": ["uint8"], "T2": ["uint8", "int8"], "T3": ["float32"], "T4": ["int32"] }, "tunables": { "WORKGROUP_SIZE": { "default": 256 }, "ATTEND_WORKGROUP": { "default": 64 }, "ATTEND_QUERY_TILE": { "default": 4 }, "ATTEND_MIN_WORKGROUPS": { "default": 1024 }, "DP4A_TILE_M": { "default": 32 }, "DP4A_TILE_N": { "default": 64 }, "DP4A_GEMV_TILE_N": { "default": 32 }, "DP4A_GEMV_K_UNROLL": { "default": 8 }, "DP4A_K_TILE_WORDS": { "default": 32 }, "EXACT_MATRIX_MIN_INVOCATIONS": { "default": 16384 } }, "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", "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32", "batchSize": "dim(shapes.input, 0)", "seqLen": "dim(shapes.input, 1)", "inputHidden": "dim(shapes.input, 2)", "hiddenX3": "dim(shapes.weight, 1)", "numHeads": "attrs.num_heads", "hidden": "hiddenX3 / 3", "headSize": "hidden / max(1, numHeads)", "headPad": "ceilDiv(headSize, 4) * 4", "headPadV4": "headPad / 4", "qkvHidden": "numHeads * headPad", "gemmRows": "batchSize * seqLen", "pastLen": "dim(shapes.past, 3) if present.past else 0", "totalLen": "pastLen + seqLen", "causal": "attrs.unidirectional != 0 and seqLen > 1", "maskRank": "ranks.mask_index if present.mask_index else 0", "maskCount": "numel(shapes.mask_index) if present.mask_index else 0", "maskKind": "\"none\" if not present.mask_index else (\"raw3d\" if maskRank == 3 else (\"raw2d\" if maskRank == 2 else (\"endpos\" if (maskRank == 1 and maskCount == batchSize) else (\"endstart\" if (maskRank == 1 and maskCount == 2 * batchSize) else \"unsupported\"))))", "maskFilter": "attrs.mask_filter_value", "maskPerQuery": "causal or maskKind == \"raw3d\"", "maskRows": "seqLen if maskPerQuery else 1", "maskElements": "batchSize * maskRows * totalLen", "maskSeqStride": "totalLen if maskPerQuery else 0", "maskBatchStride": "maskRows * totalLen", "causalLimit": "pastLen if causal else totalLen", "scaleValue": "attrs.scale if (has(attrs, \"scale\") and attrs.scale != 0) else 1 / pow(headSize, 0.5)", "weightScalePerColumn": "numel(shapes.weight_scale) == hiddenX3", "weightZeroPerColumn": "(numel(shapes.weight_zero_point) == hiddenX3) if present.weight_zero_point else false", "aScalar": "\"u32\"", "bScalar": "\"u32\" if tensorDtypes.weight == \"uint8\" else \"i32\"", "aUnsigned": "true", "bUnsigned": "tensorDtypes.weight == \"uint8\"", "aZeroPerRow": false, "intelXe3SignedProjection": "tensorDtypes.weight == \"int8\" and has(device.adapterInfo, \"vendor\") and device.adapterInfo.vendor == \"intel\" and has(device.adapterInfo, \"architecture\") and device.adapterInfo.architecture == \"xe-3lpg\"", "gemmTileM": "1 if gemmRows == 1 else tunables.DP4A_TILE_M", "gemmTileN": "tunables.DP4A_GEMV_TILE_N if gemmRows == 1 else (min(deviceWorkgroupCap, 2 * tunables.DP4A_TILE_N) if intelXe3SignedProjection else tunables.DP4A_TILE_N)", "gemmKtWords": "min(2 * tunables.DP4A_K_TILE_WORDS, floor((device.limits.maxComputeWorkgroupStorageSize - gemmTileM * 4) / (gemmTileM * 4))) if has(device.adapterInfo, \"architecture\") and device.adapterInfo.architecture == \"maxwell\" and ceilDiv(gemmRows, gemmTileM) >= 8 else tunables.DP4A_K_TILE_WORDS", "gemmStorageBytes": "gemmTileM * gemmKtWords * 4 + gemmTileM * 4", "dp4aOk": "gemmTileN <= deviceWorkgroupCap and gemmStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(hiddenX3, gemmTileN) <= DISPATCH_FOLD_WIDTH and ceilDiv(gemmRows, gemmTileM) <= DISPATCH_FOLD_WIDTH", "sgmatTileN": 64, "sgmatTileK": 32, "sgmatColumnGroups": 2, "sgmatStorageRowGroups": "floor((device.limits.maxComputeWorkgroupStorageSize - sgmatTileN * sgmatTileK * 4) / (16 * sgmatTileK * 4))", "sgmatRowGroups": "max(1, min(4, floor(deviceWorkgroupCap / (32 * sgmatColumnGroups)), sgmatStorageRowGroups, max(floor(sgmatStorageRowGroups / 2), floor(gemmRows / 16))))", "sgmatTileM": "16 * sgmatRowGroups", "sgmatSubgroups": "sgmatRowGroups * sgmatColumnGroups", "sgmatWorkgroup": "32 * sgmatSubgroups", "sgmatStagingBytes": "(sgmatTileM + sgmatTileN) * sgmatTileK * 4", "sgmatScratchBytes": "sgmatSubgroups * 4 * 64 * 4", "sgmatReuseStaging": "sgmatStagingBytes + sgmatScratchBytes > device.limits.maxComputeWorkgroupStorageSize", "sgmatSharedBytes": "sgmatStagingBytes + (0 if sgmatReuseStaging else sgmatScratchBytes)", "sgmatOk": "sgmatWorkgroup <= deviceWorkgroupCap and sgmatSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(hiddenX3, sgmatTileN) <= DISPATCH_FOLD_WIDTH and ceilDiv(gemmRows, sgmatTileM) <= DISPATCH_FOLD_WIDTH and gemmRows >= sgmatTileM and hiddenX3 >= sgmatTileN / 2 and inputHidden >= sgmatTileK and ceilDiv(gemmRows, sgmatTileM) * ceilDiv(hiddenX3, sgmatTileN) * sgmatWorkgroup >= tunables.EXACT_MATRIX_MIN_INVOCATIONS", "attendUnits": "seqLen * numHeads * batchSize", "attendNarrowTilePreferred": "variableSubgroup16To32 and seqLen >= 4 * tunables.ATTEND_WORKGROUP and headPadV4 >= 1 and headPadV4 <= 2 * device.adapterInfo.subgroupMinSize", "attendWorkgroup": "min(tunables.ATTEND_WORKGROUP, device.adapterInfo.subgroupMinSize) if device.features.has(\"subgroups\") and not device.features.has(\"shader-f16\") and seqLen >= 2 * device.adapterInfo.subgroupMinSize and attendUnits < 2 * tunables.ATTEND_MIN_WORKGROUPS else (min(tunables.ATTEND_WORKGROUP, 2 * device.adapterInfo.subgroupMinSize) if attendNarrowTilePreferred else tunables.ATTEND_WORKGROUP)", "queryTileCap": "max(1, min(tunables.ATTEND_QUERY_TILE, floor(attendWorkgroup / headPadV4), seqLen))", "queryTile": "max(1, min(queryTileCap, floor(attendUnits / max(1, tunables.ATTEND_MIN_WORKGROUPS))))", "vStripes": "max(1, floor(attendWorkgroup / (queryTile * headPadV4)))", "attendSharedBytes": "32 * queryTile * headPadV4 + 12 * queryTile * attendWorkgroup + 4 * queryTile + 16 * vStripes * queryTile * headPadV4", "attendOk": "headPadV4 >= 1 and attendSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and attendWorkgroup <= deviceWorkgroupCap and ceilDiv(seqLen, queryTile) <= DISPATCH_FOLD_WIDTH and numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and batchSize <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "shapeOk": "ranks.input == 3 and ranks.weight == 2 and ranks.bias == 1 and numHeads >= 1 and hiddenX3 >= 3 and hiddenX3 % 3 == 0 and hidden % numHeads == 0 and headSize >= 1 and inputHidden >= 1 and dim(shapes.weight, 0) == inputHidden and numel(shapes.bias) == hiddenX3 and numel(shapes.input_scale) == 1 and (numel(shapes.weight_scale) == 1 or numel(shapes.weight_scale) == hiddenX3) and sameShape(shapes.output, [batchSize, seqLen, hidden]) and seqLen >= 1 and batchSize >= 1", "maskOk": "maskKind != \"unsupported\" and (maskRank != 2 or (dim(shapes.mask_index, 0) == batchSize and dim(shapes.mask_index, 1) == totalLen)) and (maskRank != 3 or (dim(shapes.mask_index, 0) == batchSize and dim(shapes.mask_index, 1) == seqLen and dim(shapes.mask_index, 2) == totalLen))", "scratchOk": "gemmRows * hiddenX3 * 4 <= device.limits.maxStorageBufferBindingSize and 2 * batchSize * totalLen * qkvHidden * 4 <= device.limits.maxStorageBufferBindingSize and maskElements * 4 <= device.limits.maxStorageBufferBindingSize" }, "when": ["shapeOk", "attendOk", "scratchOk"], "bindings": { "a": { "arg": "input", "elementType": "$aScalar" }, "b": { "arg": "weight", "elementType": "$bScalar" }, "a_zero_point": { "arg": "input_zero_point", "elementType": "$aScalar", "length": 1 }, "b_zero_point": { "arg": "weight_zero_point", "elementType": "$bScalar" }, "y": { "scratch": "qkvAcc", "elementType": "i32" }, "gemm_params": { "name": "params", "struct": [ { "name": "M", "type": "u32", "value": "gemmRows" }, { "name": "N", "type": "u32", "value": "hiddenX3" }, { "name": "K", "type": "u32", "value": "inputHidden" }, { "name": "aBatchStride4", "type": "u32", "value": 0 }, { "name": "bBatchStride", "type": "u32", "value": 0 }, { "name": "yBatchStride", "type": "u32", "value": 0 } ] }, "sgmat_params": { "name": "params", "struct": [ { "name": "M", "type": "u32", "value": "gemmRows" }, { "name": "N", "type": "u32", "value": "hiddenX3" }, { "name": "K", "type": "u32", "value": "inputHidden" } ] }, "qkv_accumulator": { "scratch": "qkvAcc", "buffer": "read-only-storage", "elementType": "i32" }, "q_scratch": { "scratch": "qScratch", "elementType": "f32" }, "kv_scratch": { "scratch": "kvScratch", "elementType": "f32" }, "kv_scratch_read": { "scratch": "kvScratch", "name": "kv_scratch", "buffer": "read-only-storage", "elementType": "f32" }, "scatter_params": { "name": "params", "struct": [ { "name": "count", "type": "u32", "value": "gemmRows * 3 * qkvHidden" }, { "name": "seq", "type": "u32", "value": "seqLen" }, { "name": "batch", "type": "u32", "value": "batchSize" }, { "name": "total", "type": "u32", "value": "totalLen" }, { "name": "past", "type": "u32", "value": "pastLen" } ] }, "attn_mask_write": { "scratch": "attnMask", "name": "attn_mask", "elementType": "f32" }, "attn_mask": { "scratch": "attnMask", "buffer": "read-only-storage", "elementType": "f32" }, "mask_params": { "name": "params", "struct": [ { "name": "count", "type": "u32", "value": "maskElements" }, { "name": "total", "type": "u32", "value": "totalLen" }, { "name": "rows", "type": "u32", "value": "maskRows" }, { "name": "causalLimit", "type": "u32", "value": "causalLimit" } ] }, "past_params": { "name": "params", "struct": [ { "name": "count", "type": "u32", "value": "2 * batchSize * numHeads * pastLen * headSize" }, { "name": "batch", "type": "u32", "value": "batchSize" }, { "name": "total", "type": "u32", "value": "totalLen" }, { "name": "length", "type": "u32", "value": "pastLen" } ] }, "present_params": { "name": "params", "struct": [ { "name": "count", "type": "u32", "value": "2 * batchSize * numHeads * totalLen * headSize" }, { "name": "batch", "type": "u32", "value": "batchSize" }, { "name": "total", "type": "u32", "value": "totalLen" }, { "name": "length", "type": "u32", "value": "totalLen" } ] }, "query": { "scratch": "qScratch", "buffer": "read-only-storage", "elementType": "vec4" }, "key_value": { "scratch": "kvScratch", "buffer": "read-only-storage", "elementType": "vec4" }, "attend_params": { "name": "params", "struct": [ { "name": "qSeq", "type": "u32", "value": "seqLen" }, { "name": "kvSeq", "type": "u32", "value": "totalLen" }, { "name": "batch", "type": "u32", "value": "batchSize" }, { "name": "maskBatchStride", "type": "u32", "value": "maskBatchStride" }, { "name": "maskSeqStride", "type": "u32", "value": "maskSeqStride" }, { "name": "scale", "type": "f32", "value": "scaleValue" } ] } }, "variants": [ { "id": "sgmat_plain_mask_izp_wzp", "priority": 20, "when": ["not present.past", "not present.present", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "packed_plain_mask_izp_wzp", "priority": 10, "when": ["not present.past", "not present.present", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "sgmat_plain_mask_izp_nowzp", "priority": 20, "when": ["not present.past", "not present.present", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "packed_plain_mask_izp_nowzp", "priority": 10, "when": ["not present.past", "not present.present", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "sgmat_plain_mask_noizp_wzp", "priority": 20, "when": ["not present.past", "not present.present", "present.mask_index and maskOk", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "packed_plain_mask_noizp_wzp", "priority": 10, "when": ["not present.past", "not present.present", "present.mask_index and maskOk", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "sgmat_plain_mask_noizp_nowzp", "priority": 20, "when": ["not present.past", "not present.present", "present.mask_index and maskOk", "not present.input_zero_point", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "packed_plain_mask_noizp_nowzp", "priority": 10, "when": ["not present.past", "not present.present", "present.mask_index and maskOk", "not present.input_zero_point", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "sgmat_plain_nomask_izp_wzp", "priority": 20, "when": ["not present.past", "not present.present", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "packed_plain_nomask_izp_wzp", "priority": 10, "when": ["not present.past", "not present.present", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "sgmat_plain_nomask_izp_nowzp", "priority": 20, "when": ["not present.past", "not present.present", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "packed_plain_nomask_izp_nowzp", "priority": 10, "when": ["not present.past", "not present.present", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "sgmat_plain_nomask_noizp_wzp", "priority": 20, "when": ["not present.past", "not present.present", "not present.mask_index", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "packed_plain_nomask_noizp_wzp", "priority": 10, "when": ["not present.past", "not present.present", "not present.mask_index", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "sgmat_plain_nomask_noizp_nowzp", "priority": 20, "when": ["not present.past", "not present.present", "not present.mask_index", "not present.input_zero_point", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "packed_plain_nomask_noizp_nowzp", "priority": 10, "when": ["not present.past", "not present.present", "not present.mask_index", "not present.input_zero_point", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } } ] }, { "id": "sgmat_present_mask_izp_wzp", "priority": 20, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_present_mask_izp_wzp", "priority": 10, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_present_mask_izp_nowzp", "priority": 20, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_present_mask_izp_nowzp", "priority": 10, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_present_mask_noizp_wzp", "priority": 20, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_present_mask_noizp_wzp", "priority": 10, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_present_mask_noizp_nowzp", "priority": 20, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "not present.input_zero_point", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_present_mask_noizp_nowzp", "priority": 10, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "not present.input_zero_point", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_present_nomask_izp_wzp", "priority": 20, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_present_nomask_izp_wzp", "priority": 10, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_present_nomask_izp_nowzp", "priority": 20, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_present_nomask_izp_nowzp", "priority": 10, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_present_nomask_noizp_wzp", "priority": 20, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_present_nomask_noizp_wzp", "priority": 10, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_present_nomask_noizp_nowzp", "priority": 20, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "not present.input_zero_point", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_present_nomask_noizp_nowzp", "priority": 10, "when": ["not present.past", "present.present", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "not present.input_zero_point", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_past_mask_izp_wzp", "priority": 20, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_past_mask_izp_wzp", "priority": 10, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_past_mask_izp_nowzp", "priority": 20, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_past_mask_izp_nowzp", "priority": 10, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_past_mask_noizp_wzp", "priority": 20, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_past_mask_noizp_wzp", "priority": 10, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_past_mask_noizp_nowzp", "priority": 20, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "not present.input_zero_point", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_past_mask_noizp_nowzp", "priority": 10, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "present.mask_index and maskOk", "not present.input_zero_point", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["mask_index", "attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_past_nomask_izp_wzp", "priority": 20, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_past_nomask_izp_wzp", "priority": 10, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_past_nomask_izp_nowzp", "priority": 20, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_past_nomask_izp_nowzp", "priority": 10, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "present.input_zero_point", "numel(shapes.input_zero_point) == 1", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "a_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_past_nomask_noizp_wzp", "priority": 20, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_past_nomask_noizp_wzp", "priority": 10, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "not present.input_zero_point", "present.weight_zero_point", "numel(shapes.weight_zero_point) == 1 or numel(shapes.weight_zero_point) == hiddenX3", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "b_zero_point", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "sgmat_past_nomask_noizp_nowzp", "priority": 20, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "not present.input_zero_point", "not present.weight_zero_point", "sgmatOk", "wave32Effective"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectExactMatrix", "shader": "quant-exact-matrix.wgsl.jinja", "bindings": ["a", "b", "y", "sgmat_params"], "dispatch": { "x": "ceilDiv(hiddenX3, sgmatTileN)", "y": "ceilDiv(gemmRows, sgmatTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] }, { "id": "packed_past_nomask_noizp_nowzp", "priority": 10, "when": ["present.past", "present.present", "ranks.past == 5", "dim(shapes.past, 0) == 2 and dim(shapes.past, 1) == batchSize", "dim(shapes.past, 2) == numHeads and dim(shapes.past, 4) == headSize", "sameShape(shapes.present, [2, batchSize, numHeads, totalLen, headSize])", "not present.mask_index", "not present.input_zero_point", "not present.weight_zero_point", "dp4aOk"], "derive": { "hasAZero": "present.input_zero_point", "hasBZero": "present.weight_zero_point", "bZeroPerColumn": "weightZeroPerColumn", "tileM": "gemmTileM", "tileN": "gemmTileN", "ktWords": "gemmKtWords", "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" }, "intermediates": [ { "id": "qkvAcc", "dtype": "int32", "shape": "[gemmRows * hiddenX3]" }, { "id": "qScratch", "dtype": "float32", "shape": "[batchSize * seqLen * qkvHidden]" }, { "id": "kvScratch", "dtype": "float32", "shape": "[2 * batchSize * totalLen * qkvHidden]" }, { "id": "attnMask", "dtype": "float32", "shape": "[maskElements]" } ], "passes": [ { "id": "past_state", "name": "QAttention.PastState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"past_in\"" }, "bindings": ["past", "kv_scratch", "past_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * pastLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "mask", "name": "QAttention.Mask", "shader": "qattention-mask.wgsl.jinja", "bindings": ["attn_mask_write", "mask_params"], "dispatch": { "x": "min(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((maskElements), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "project", "name": "QAttention.ProjectPackedInteger", "shader": "quant-dp4a-matmul.wgsl.jinja", "bindings": ["a", "b", "y", "gemm_params"], "dispatch": { "x": "ceilDiv(hiddenX3, gemmTileN)", "y": "ceilDiv(gemmRows, gemmTileM)" } }, { "id": "scatter", "name": "QAttention.Dequantize", "shader": "qattention-qkv-scatter.wgsl.jinja", "bindings": ["qkv_accumulator", "bias", "input_scale", "weight_scale", "q_scratch", "kv_scratch", "scatter_params"], "dispatch": { "x": "min(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((gemmRows * 3 * qkvHidden), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "attend", "name": "QAttention.Attend", "shader": "qattention-attend.wgsl.jinja", "bindings": ["query", "key_value", "attn_mask", "output", "attend_params"], "dispatch": { "x": "ceilDiv(seqLen, queryTile)", "y": "numHeads", "z": "batchSize" } }, { "id": "present_state", "name": "QAttention.PresentState", "shader": "qattention-kv-permute.wgsl.jinja", "derive": { "direction": "\"present_out\"" }, "bindings": ["kv_scratch_read", "present", "present_params"], "dispatch": { "x": "min(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((2 * batchSize * numHeads * totalLen * headSize), (tunables.WORKGROUP_SIZE)), 65535)", "z": 1 } } ] } ] }