{ "domain": "com.microsoft", "name": "PackedAttention", "sinceVersion": 1, "inputs": { "inputT": { "onnx": "input", "dtype": "T", "rank": 2 }, "weightsT": { "onnx": "weights", "dtype": "T", "rank": 2 }, "biasT": { "onnx": "bias", "dtype": "T", "rank": 1 }, "tokenOffsetT": { "onnx": "token_offset", "dtype": "M", "rank": 2, "storage": "int32" }, "cumulativeSequenceLengthT": { "onnx": "cumulative_sequence_length", "dtype": "M", "rank": 1, "storage": "int32" }, "attentionBiasT": { "onnx": "attention_bias", "dtype": "T", "rank": 4, "optional": true } }, "outputs": { "outputT": { "onnx": "output", "dtype": "T", "rank": 2, "shape": "[tokenCount, vHidden]" } }, "attributes": { "num_heads": {}, "qkv_hidden_sizes": {}, "scale": {} }, "attributeConstraints": { "num_heads": { "required": true } }, "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] }, "tunables": { "CLUSTER_TILE_K": { "default": 16 }, "SCALAR_WORKGROUP_SIZE": { "default": 64 }, "MATRIX_DIRECT_MIN_GROUPS": { "default": 128 } }, "derive": { "numHeads": "attrs.num_heads", "qNumHeads": "numHeads", "tokenCount": "dim(shapes.inputT, 0) if ranks.inputT == 2 else 0", "inputHidden": "dim(shapes.inputT, 1) if ranks.inputT == 2 else 0", "projection": "dim(shapes.weightsT, 1) if ranks.weightsT == 2 else 0", "qkvSizesGiven": "has(attrs, \"qkv_hidden_sizes\")", "qHidden": "attrs.qkv_hidden_sizes[0] if qkvSizesGiven else projection / 3", "kHidden": "attrs.qkv_hidden_sizes[1] if qkvSizesGiven else qHidden", "vHidden": "attrs.qkv_hidden_sizes[2] if qkvSizesGiven else qHidden", "headDim": "qHidden / numHeads if numHeads > 0 and qHidden % numHeads == 0 else 0", "vHeadDim": "vHidden / numHeads if numHeads > 0 and vHidden % numHeads == 0 else 0", "batchSize": "dim(shapes.tokenOffsetT, 0) if ranks.tokenOffsetT == 2 else 0", "seqPad": "dim(shapes.tokenOffsetT, 1) if ranks.tokenOffsetT == 2 else 0", "attentionScale": "attrs.scale if has(attrs, \"scale\") else 0", "qRowStride": "projection", "qHeadStride": "headDim", "kRowStride": "projection", "kHeadStride": "headDim", "kOffset": "qHidden", "vRowStride": "projection", "vHeadStride": "vHeadDim", "vOffset": "qHidden + kHidden", "oRowStride": "vHidden", "oHeadStride": "vHeadDim", "attnBiasHeads": "dim(shapes.attentionBiasT, 1) if present.attentionBiasT and ranks.attentionBiasT == 4 else 0", "maskSeqStride": "seqPad", "maskHeadStride": "0 if attnBiasHeads == 1 else seqPad * seqPad", "maskBatchStride": "0 if (not present.attentionBiasT) or ranks.attentionBiasT != 4 or dim(shapes.attentionBiasT, 0) == 1 else attnBiasHeads * seqPad * seqPad", "dtypeOk": "tensorDtypes.weightsT == tensorDtypes.inputT and tensorDtypes.biasT == tensorDtypes.inputT and tensorDtypes.outputT == tensorDtypes.inputT and f16Ok(dtypes.T)", "sizesOk": "qHidden == kHidden and qHidden + kHidden + vHidden == projection and headDim > 0 and vHeadDim > 0", "weightsOk": "ranks.weightsT == 2 and dim(shapes.weightsT, 0) == inputHidden and ranks.biasT == 1 and dim(shapes.biasT, 0) == projection", "scheduleOk": "ranks.tokenOffsetT == 2 and ranks.cumulativeSequenceLengthT == 1 and dim(shapes.cumulativeSequenceLengthT, 0) == batchSize + 1 and batchSize > 0 and seqPad > 0 and numel(shapes.tokenOffsetT) >= tokenCount", "outputShapeOk": "ranks.outputT == 2 and dim(shapes.outputT, 0) == tokenCount and dim(shapes.outputT, 1) == vHidden", "attnBiasOk": "present.attentionBiasT and ranks.attentionBiasT == 4 and (dim(shapes.attentionBiasT, 0) == 1 or dim(shapes.attentionBiasT, 0) == batchSize) and (attnBiasHeads == 1 or attnBiasHeads == numHeads) and dim(shapes.attentionBiasT, 2) == seqPad and dim(shapes.attentionBiasT, 3) == seqPad and tensorDtypes.attentionBiasT == tensorDtypes.inputT", "hasAttnBias": "present.attentionBiasT", "baseContract": "dtypeOk and weightsOk and sizesOk and scheduleOk and outputShapeOk and numHeads > 0 and tokenCount > 0 and inputHidden > 0", "usesF16": "tensorDtypes.inputT == \"float16\"", "scalar": "\"f16\" if usesF16 else \"f32\"", "outputScalar": "scalar", "inputVec4": "\"vec4\" if usesF16 else \"vec4\"", "qkvBytes": "2 if usesF16 else 4", "packedStageBytes": "qkvBytes", "projectionBytes": "2 * 64 * 16 * qkvBytes", "projectionOk": "projectionBytes <= device.limits.maxComputeWorkgroupStorageSize and 256 <= device.limits.maxComputeInvocationsPerWorkgroup and 16 <= device.limits.maxComputeWorkgroupSizeX and 16 <= device.limits.maxComputeWorkgroupSizeY", "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", "directProjectionRows": 64, "directProjectionCols": 32, "directProjectionSharedBytes": "2048 * qkvBytes + 8192", "directProjectionOk": "wave32Effective and inputHidden % 32 == 0 and projection % directProjectionCols == 0 and 128 <= deviceWorkgroupCap and directProjectionSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(projection, directProjectionCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(tokenCount, directProjectionRows), DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "stagedProjectionRows": 32, "stagedProjectionCols": 32, "stagedProjectionSharedBytes": "2048 * qkvBytes + 4096", "stagedProjectionOk": "wave32Effective and inputHidden % 32 == 0 and projection % stagedProjectionCols == 0 and 128 <= deviceWorkgroupCap and stagedProjectionSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(projection, stagedProjectionCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "directProjectionFullGroups": "floor(tokenCount / directProjectionRows) * ceilDiv(projection, directProjectionCols)", "directProjectionEffectiveGroups": "directProjectionFullGroups * inputHidden / (inputHidden + 32)", "clusterWidthAligned": "headDim % 16 == 0 and vHeadDim % 4 == 0 and projection % 4 == 0", "headDimV4": "headDim / 4", "qRowStrideV4": "qRowStride / 4", "qHeadStrideV4": "qHeadStride / 4", "kRowStrideV4": "kRowStride / 4", "kHeadStrideV4": "kHeadStride / 4", "kOffsetV4": "kOffset / 4", "vRowStrideV4": "vRowStride / 4", "vHeadStrideV4": "vHeadStride / 4", "vOffsetV4": "vOffset / 4", "oRowStrideV4": "oRowStride / 4", "oHeadStrideV4": "oHeadStride / 4", "scalarWg": "tunables.SCALAR_WORKGROUP_SIZE", "scalarBytes": "4 * scalarWg + 4 * vHeadDim + 12", "scalarOk": "scalarWg <= device.limits.maxComputeInvocationsPerWorkgroup and scalarWg <= device.limits.maxComputeWorkgroupSizeX and scalarBytes <= device.limits.maxComputeWorkgroupStorageSize and numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "clusterTileQ": "min(16, pow2ceil(tokenCount)) if vHeadDim != headDim else 16", "clusterBaseLpq": "8 if headDim % 32 == 0 else 4", "clusterTileK": "min(tunables.CLUSTER_TILE_K, pow2ceil(tokenCount)) if vHeadDim != headDim else tunables.CLUSTER_TILE_K", "clusterTileBytes": "clusterTileK * (headDim + vHeadDim) * packedStageBytes", "clusterLaneTarget": "min(32, max(clusterBaseLpq, pow2ceil(headDim / 16)), device.adapterInfo.subgroupMinSize, pow2ceil(floor(min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX) / clusterTileQ) + 1) / 2) if vHeadDim == headDim and clusterWidthAligned and clusterTileBytes <= device.limits.maxComputeWorkgroupStorageSize and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize >= clusterBaseLpq else clusterBaseLpq", "clusterLpq": "clusterLaneTarget if ceilDiv(headDim, 4 * clusterBaseLpq) >= 2 * ceilDiv(headDim, 4 * clusterLaneTarget) else clusterBaseLpq", "clusterWg": "clusterTileQ * clusterLpq", "clusterNosgBytes": "clusterTileBytes + 4 * clusterWg", "clusterTiles": "ceilDiv(tokenCount, clusterTileQ)", "clusterShapeOk": "clusterWidthAligned and clusterWg <= device.limits.maxComputeInvocationsPerWorkgroup and clusterWg <= device.limits.maxComputeWorkgroupSizeX and numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "clusterOk": "clusterShapeOk and clusterTileBytes <= device.limits.maxComputeWorkgroupStorageSize", "clusterNosgOk": "clusterShapeOk and clusterNosgBytes <= device.limits.maxComputeWorkgroupStorageSize", "clusterSubgroupsOk": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize >= clusterLpq and device.adapterInfo.subgroupMinSize % clusterLpq == 0 and device.adapterInfo.subgroupMaxSize >= device.adapterInfo.subgroupMinSize and device.adapterInfo.subgroupMaxSize % clusterLpq == 0", "portableLaneTarget": "min(64, pow2ceil(max(headDimV4, vHeadDim / 4) / 4), device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", "portableLpq": "max(clusterLpq, pow2ceil(portableLaneTarget + 1) / 2) if vHeadDim != headDim else (64 if portableLaneTarget >= 64 and headDimV4 % 64 == 0 else (32 if portableLaneTarget >= 32 and headDimV4 % 32 == 0 else (16 if portableLaneTarget >= 16 and headDimV4 % 16 == 0 else clusterLpq)))", "portableTileQ": "min(clusterTileQ, pow2ceil(tokenCount), max(1, floor(min(clusterWg, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX) / portableLpq)))", "portableWg": "portableTileQ * portableLpq", "portableKeyCapacity": "max(1, floor(device.limits.maxComputeWorkgroupStorageSize / ((headDim + vHeadDim) * packedStageBytes + 4 * portableWg)))", "portableTileK": "min(clusterTileK, pow2ceil(tokenCount), pow2ceil(portableKeyCapacity + 1) / 2)", "clusterDistributedBytes": "portableTileK * (headDim + vHeadDim) * packedStageBytes + 4 * max(portableWg * portableTileK, 2 * portableWg + portableTileQ * portableTileK)", "clusterDistributedOk": "clusterDistributedBytes <= device.limits.maxComputeWorkgroupStorageSize", "clusterUseDistributed": "clusterDistributedOk and (clusterLpq > 4 or vHeadDim != headDim)", "portableTiles": "ceilDiv(tokenCount, portableTileQ if clusterUseDistributed else clusterTileQ)", "clusterPortableOk": "clusterShapeOk and (clusterUseDistributed or clusterNosgOk) and portableWg <= device.limits.maxComputeInvocationsPerWorkgroup and portableWg <= device.limits.maxComputeWorkgroupSizeX and ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "clusterPreferPortable": "vHeadDim != headDim and clusterPortableOk and clusterUseDistributed and portableLpq > clusterLpq", "clusterReusePoor": "vHeadDim != headDim and tokenCount < clusterLpq", "portableProjectionRows": "32 if not usesF16 and not device.features.has(\"subgroups\") and ceilDiv(tokenCount, 32) * ceilDiv(projection, 32) >= 2 * deviceWorkgroupCap else 64" }, "bindings": { "projInput": { "arg": "inputT", "name": "a" }, "projWeights": { "arg": "weightsT", "name": "b" }, "projBias": { "arg": "biasT", "name": "bias" }, "projOut": { "scratch": "qkvProj", "name": "y", "elementType": "$T" }, "gemmParams": { "name": "params", "struct": [ { "name": "M", "type": "u32", "value": "tokenCount" }, { "name": "N", "type": "u32", "value": "projection" }, { "name": "K", "type": "u32", "value": "inputHidden" }, { "name": "alpha", "type": "f32", "value": 1 }, { "name": "beta", "type": "f32", "value": 1 } ] }, "qkv": { "scratch": "qkvProj", "name": "query", "buffer": "read-only-storage", "elementType": "$T" }, "qkv4": { "scratch": "qkvProj", "name": "query", "buffer": "read-only-storage", "elementType": "$inputVec4" }, "tokenOffset": { "arg": "tokenOffsetT", "name": "token_offset" }, "cumSeqLen": { "arg": "cumulativeSequenceLengthT", "name": "cumulative_sequence_length" }, "attnBias": { "arg": "attentionBiasT", "name": "attn_bias" }, "output": { "arg": "outputT", "name": "output" }, "output4": { "arg": "outputT", "name": "output", "elementType": "$inputVec4" }, "params": { "name": "params", "struct": [ { "name": "tokenCount", "type": "u32", "value": "tokenCount" }, { "name": "seqPad", "type": "u32", "value": "seqPad" }, { "name": "scale", "type": "f32", "value": "attentionScale" } ] }, "paramsMask": { "name": "params", "struct": [ { "name": "tokenCount", "type": "u32", "value": "tokenCount" }, { "name": "seqPad", "type": "u32", "value": "seqPad" }, { "name": "scale", "type": "f32", "value": "attentionScale" }, { "name": "maskBatchStride", "type": "u32", "value": "maskBatchStride" }, { "name": "maskHeadStride", "type": "u32", "value": "maskHeadStride" }, { "name": "maskSeqStride", "type": "u32", "value": "maskSeqStride" } ] }, "matrixInput": { "arg": "inputT", "name": "x" }, "matrixWeights": { "arg": "weightsT", "name": "w" }, "matrixParams": { "name": "params", "struct": [{ "name": "M", "type": "u32", "value": "tokenCount" }] } }, "variants": [ { "id": "matrix_direct_f32_flash", "priority": 140, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "not present.attentionBiasT", "directProjectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK", "distributedSoftmax": false }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)", "clusterReusePoor or clusterPreferPortable"] }, { "id": "matrix_staged_f32_flash", "priority": 120, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "not present.attentionBiasT", "stagedProjectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK", "distributedSoftmax": false }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["clusterReusePoor or clusterPreferPortable"] }, { "id": "f32_flash", "priority": 40, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "not present.attentionBiasT", "projectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["clusterReusePoor or clusterPreferPortable"] }, { "id": "matrix_direct_f32_flash_nosg", "priority": 139, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "not present.attentionBiasT", "directProjectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)", "clusterReusePoor or (clusterOk and clusterSubgroupsOk and not clusterPreferPortable)"] }, { "id": "matrix_staged_f32_flash_nosg", "priority": 119, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "not present.attentionBiasT", "stagedProjectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["clusterReusePoor or (clusterOk and clusterSubgroupsOk and not clusterPreferPortable)"] }, { "id": "f32_flash_nosg", "priority": 20, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "not present.attentionBiasT", "projectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["clusterReusePoor"] }, { "id": "matrix_direct_f32_serial", "priority": 100, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "not present.attentionBiasT", "directProjectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "output", "params"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)"] }, { "id": "matrix_staged_f32_serial", "priority": 80, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "not present.attentionBiasT", "stagedProjectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "output", "params"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] } }, { "id": "f32_serial", "priority": 0, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "not present.attentionBiasT", "projectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "output", "params"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] }, { "id": "matrix_direct_f32_attn_bias_flash", "priority": 140, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "attnBiasOk", "directProjectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK", "distributedSoftmax": false }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)", "clusterReusePoor or clusterPreferPortable"] }, { "id": "matrix_staged_f32_attn_bias_flash", "priority": 120, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "attnBiasOk", "stagedProjectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK", "distributedSoftmax": false }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["clusterReusePoor or clusterPreferPortable"] }, { "id": "f32_attn_bias_flash", "priority": 40, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "attnBiasOk", "projectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["clusterReusePoor or clusterPreferPortable"] }, { "id": "matrix_direct_f32_attn_bias_flash_nosg", "priority": 139, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "attnBiasOk", "directProjectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)", "clusterReusePoor or (clusterOk and clusterSubgroupsOk and not clusterPreferPortable)"] }, { "id": "matrix_staged_f32_attn_bias_flash_nosg", "priority": 119, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "attnBiasOk", "stagedProjectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["clusterReusePoor or (clusterOk and clusterSubgroupsOk and not clusterPreferPortable)"] }, { "id": "f32_attn_bias_flash_nosg", "priority": 20, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "attnBiasOk", "projectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["clusterReusePoor"] }, { "id": "matrix_direct_f32_attn_bias_serial", "priority": 100, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "attnBiasOk", "directProjectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)"] }, { "id": "matrix_staged_f32_attn_bias_serial", "priority": 80, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "attnBiasOk", "stagedProjectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] } }, { "id": "f32_attn_bias_serial", "priority": 0, "when": ["baseContract", "tensorDtypes.inputT == \"float32\"", "attnBiasOk", "projectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float32", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] }, { "id": "matrix_direct_f16_flash", "priority": 140, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "not present.attentionBiasT", "directProjectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK", "distributedSoftmax": false }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)", "clusterReusePoor or clusterPreferPortable"] }, { "id": "matrix_staged_f16_flash", "priority": 120, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "not present.attentionBiasT", "stagedProjectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK", "distributedSoftmax": false }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["clusterReusePoor or clusterPreferPortable"] }, { "id": "f16_flash", "priority": 40, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "not present.attentionBiasT", "projectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["clusterReusePoor or clusterPreferPortable"] }, { "id": "matrix_direct_f16_flash_nosg", "priority": 139, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "not present.attentionBiasT", "directProjectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)", "clusterReusePoor or (clusterOk and clusterSubgroupsOk and not clusterPreferPortable)"] }, { "id": "matrix_staged_f16_flash_nosg", "priority": 119, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "not present.attentionBiasT", "stagedProjectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["clusterReusePoor or (clusterOk and clusterSubgroupsOk and not clusterPreferPortable)"] }, { "id": "f16_flash_nosg", "priority": 20, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "not present.attentionBiasT", "projectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["clusterReusePoor"] }, { "id": "matrix_direct_f16_serial", "priority": 100, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "not present.attentionBiasT", "directProjectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "output", "params"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)"] }, { "id": "matrix_staged_f16_serial", "priority": 80, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "not present.attentionBiasT", "stagedProjectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "output", "params"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] } }, { "id": "f16_serial", "priority": 0, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "not present.attentionBiasT", "projectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "output", "params"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] }, { "id": "matrix_direct_f16_attn_bias_flash", "priority": 140, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "attnBiasOk", "directProjectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK", "distributedSoftmax": false }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)", "clusterReusePoor or clusterPreferPortable"] }, { "id": "matrix_staged_f16_attn_bias_flash", "priority": 120, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "attnBiasOk", "stagedProjectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK", "distributedSoftmax": false }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["clusterReusePoor or clusterPreferPortable"] }, { "id": "f16_attn_bias_flash", "priority": 40, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "attnBiasOk", "projectionOk", "clusterOk", "clusterSubgroupsOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["clusterReusePoor or clusterPreferPortable"] }, { "id": "matrix_direct_f16_attn_bias_flash_nosg", "priority": 139, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "attnBiasOk", "directProjectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)", "clusterReusePoor or (clusterOk and clusterSubgroupsOk and not clusterPreferPortable)"] }, { "id": "matrix_staged_f16_attn_bias_flash_nosg", "priority": 119, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "attnBiasOk", "stagedProjectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["clusterReusePoor or (clusterOk and clusterSubgroupsOk and not clusterPreferPortable)"] }, { "id": "f16_attn_bias_flash_nosg", "priority": 20, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "attnBiasOk", "projectionOk", "clusterPortableOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "derive": { "useSubgroups": false, "LPQ": "portableLpq if clusterUseDistributed else clusterLpq", "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ", "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK", "distributedSoftmax": "clusterUseDistributed" }, "bindings": ["qkv4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["clusterReusePoor"] }, { "id": "matrix_direct_f16_attn_bias_serial", "priority": 100, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "attnBiasOk", "directProjectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "full_projection", "repeat": { "index": "fullProjectionIndex", "count": "1 if tokenCount >= directProjectionRows else 0" }, "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "min(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(floor(tokenCount / directProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "directProjectionRows", "tileN": "directProjectionCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": "not usesF16" } } ] }, { "id": "tail_projection", "repeat": { "index": "tailProjectionIndex", "count": "1 if tokenCount % directProjectionRows != 0 else 0" }, "passes": [ { "id": "project_tail", "name": "PackedAttention.ProjectionTail", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, directProjectionCols)", "y": "ceilDiv(tokenCount % directProjectionRows, 32)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": 32, "tileN": "directProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false, "rowTileBase": "floor(tokenCount / directProjectionRows) * directProjectionRows / 32" } } ] }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "demoteWhen": ["directProjectionEffectiveGroups < tunables.MATRIX_DIRECT_MIN_GROUPS * (2 if tokenCount % directProjectionRows != 0 else 1)"] }, { "id": "matrix_staged_f16_attn_bias_serial", "priority": 80, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "attnBiasOk", "stagedProjectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-subgroup-matrix.wgsl.jinja", "bindings": ["matrixInput", "matrixWeights", "projBias", "projOut", "matrixParams"], "dispatch": { "x": "ceilDiv(projection, stagedProjectionCols)", "y": "min(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)", "z": "ceilDiv(ceilDiv(tokenCount, stagedProjectionRows), DISPATCH_FOLD_WIDTH)" }, "derive": { "yScalar": "scalar", "mmaScalar": "scalar", "inFeatures": "inputHidden", "outFeatures": "projection", "tileM": "stagedProjectionRows", "tileN": "stagedProjectionCols", "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false } }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] } }, { "id": "f16_attn_bias_serial", "priority": 0, "when": ["baseContract", "tensorDtypes.inputT == \"float16\"", "attnBiasOk", "projectionOk", "scalarOk"], "intermediates": [{ "id": "qkvProj", "dtype": "float16", "shape": "[tokenCount * projection]" }], "passes": [ { "id": "project", "name": "PackedAttention.Projection", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "derive": { "gemmTileRows": "portableProjectionRows", "gemmTileCols": "portableProjectionRows", "gemmMicroRows": "portableProjectionRows / 16", "gemmMicroCols": "portableProjectionRows / 16" }, "bindings": ["projInput", "projWeights", "projBias", "projOut", "gemmParams"], "dispatch": { "x": "min(ceilDiv(projection, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(tokenCount, portableProjectionRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } }, { "id": "attend", "name": "PackedAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "derive": { "workgroupSize": "scalarWg" }, "bindings": ["qkv", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] } ] }