{ "domain": "com.microsoft", "name": "PackedMultiHeadAttention", "sinceVersion": 1, "inputs": { "queryT": { "onnx": "query", "dtype": "T" }, "keyT": { "onnx": "key", "dtype": "T", "rank": 2, "optional": true }, "valueT": { "onnx": "value", "dtype": "T", "rank": 2, "optional": true }, "biasT": { "onnx": "bias", "dtype": "T", "rank": 1, "optional": true }, "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, vHiddenSize]" } }, "attributes": { "num_heads": {}, "scale": {} }, "attributeConstraints": { "num_heads": { "required": true } }, "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] }, "tunables": { "CLUSTER_TILE_K": { "default": 16 }, "SCALAR_WORKGROUP_SIZE": { "default": 64 } }, "derive": { "numHeads": "attrs.num_heads", "qNumHeads": "numHeads", "packedQkv": "ranks.queryT == 4", "tokenCount": "dim(shapes.queryT, 0)", "batchSize": "dim(shapes.tokenOffsetT, 0) if ranks.tokenOffsetT == 2 else 0", "seqPad": "dim(shapes.tokenOffsetT, 1) if ranks.tokenOffsetT == 2 else 0", "packedHeadDim": "dim(shapes.queryT, 3) if packedQkv else 0", "hiddenSize": "numHeads * packedHeadDim if packedQkv else (dim(shapes.queryT, 1) if ranks.queryT == 2 else 0)", "vHiddenSize": "hiddenSize if packedQkv else (dim(shapes.valueT, 1) if present.valueT and ranks.valueT == 2 else 0)", "headDim": "hiddenSize / numHeads if numHeads > 0 and hiddenSize % numHeads == 0 else 0", "vHeadDim": "vHiddenSize / numHeads if numHeads > 0 and vHiddenSize % numHeads == 0 else 0", "attentionScale": "attrs.scale if has(attrs, \"scale\") else 0", "qRowStride": "numHeads * 3 * headDim if packedQkv else hiddenSize", "qHeadStride": "3 * headDim if packedQkv else headDim", "kRowStride": "qRowStride if packedQkv else hiddenSize", "kHeadStride": "qHeadStride if packedQkv else headDim", "kOffset": "headDim if packedQkv else 0", "vRowStride": "qRowStride if packedQkv else vHiddenSize", "vHeadStride": "qHeadStride if packedQkv else vHeadDim", "vOffset": "2 * headDim if packedQkv else 0", "oRowStride": "vHiddenSize", "oHeadStride": "vHeadDim", "biasKOffset": "hiddenSize", "biasVOffset": "2 * hiddenSize", "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.outputT == tensorDtypes.queryT and (tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(dtypes.T)", "packedShapeOk": "packedQkv and not present.keyT and not present.valueT and dim(shapes.queryT, 1) == numHeads and dim(shapes.queryT, 2) == 3 and packedHeadDim > 0", "separateShapeOk": "ranks.queryT == 2 and present.keyT and present.valueT and ranks.keyT == 2 and ranks.valueT == 2 and sameShape(shapes.keyT, shapes.queryT) and dim(shapes.valueT, 0) == tokenCount and tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT", "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) == vHiddenSize", "biasOk": "present.biasT and ranks.biasT == 1 and dim(shapes.biasT, 0) == 2 * hiddenSize + vHiddenSize and tensorDtypes.biasT == tensorDtypes.queryT", "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.queryT", "hasBias": "present.biasT", "hasAttnBias": "present.attentionBiasT", "baseContract": "dtypeOk and scheduleOk and outputShapeOk and numHeads > 0 and headDim > 0 and vHeadDim > 0 and tokenCount > 0 and (packedShapeOk or separateShapeOk)", "usesF16": "tensorDtypes.queryT == \"float16\"", "scalar": "\"f16\" if usesF16 else \"f32\"", "outputScalar": "scalar", "inputVec4": "\"vec4\" if usesF16 else \"vec4\"", "qkvBytes": "2 if usesF16 else 4", "packedStageBytes": "4 if usesF16 and present.biasT and vHeadDim != headDim else qkvBytes", "clusterWidthAligned": "headDim % 16 == 0 and vHeadDim % 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)" }, "bindings": { "query": { "arg": "queryT", "name": "query" }, "key": { "arg": "keyT", "name": "key" }, "value": { "arg": "valueT", "name": "value" }, "bias": { "arg": "biasT", "name": "bias" }, "tokenOffset": { "arg": "tokenOffsetT", "name": "token_offset" }, "cumSeqLen": { "arg": "cumulativeSequenceLengthT", "name": "cumulative_sequence_length" }, "attnBias": { "arg": "attentionBiasT", "name": "attn_bias" }, "output": { "arg": "outputT", "name": "output" }, "query4": { "arg": "queryT", "name": "query", "elementType": "$inputVec4" }, "key4": { "arg": "keyT", "name": "key", "elementType": "$inputVec4" }, "value4": { "arg": "valueT", "name": "value", "elementType": "$inputVec4" }, "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" } ] } }, "variants": [ { "id": "packed_flash", "priority": 40, "when": ["baseContract", "packedShapeOk", "not present.biasT", "not present.attentionBiasT", "clusterOk", "clusterSubgroupsOk"], "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"] }, { "id": "packed_flash_nosg", "priority": 20, "when": ["baseContract", "packedShapeOk", "not present.biasT", "not present.attentionBiasT", "clusterPortableOk"], "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" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"] }, { "id": "packed_serial", "priority": 0, "when": ["baseContract", "packedShapeOk", "not present.biasT", "not present.attentionBiasT", "scalarOk"], "derive": { "workgroupSize": "scalarWg" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "bindings": ["query", "tokenOffset", "cumSeqLen", "output", "params"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] }, { "id": "packed_attn_bias_flash", "priority": 40, "when": ["baseContract", "packedShapeOk", "not present.biasT", "attnBiasOk", "clusterOk", "clusterSubgroupsOk"], "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"] }, { "id": "packed_attn_bias_flash_nosg", "priority": 20, "when": ["baseContract", "packedShapeOk", "not present.biasT", "attnBiasOk", "clusterPortableOk"], "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" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"] }, { "id": "packed_attn_bias_serial", "priority": 0, "when": ["baseContract", "packedShapeOk", "not present.biasT", "attnBiasOk", "scalarOk"], "derive": { "workgroupSize": "scalarWg" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "bindings": ["query", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] }, { "id": "packed_bias_flash", "priority": 40, "when": ["baseContract", "packedShapeOk", "biasOk", "not present.attentionBiasT", "clusterOk", "clusterSubgroupsOk"], "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "bias", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"] }, { "id": "packed_bias_flash_nosg", "priority": 20, "when": ["baseContract", "packedShapeOk", "biasOk", "not present.attentionBiasT", "clusterPortableOk"], "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" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "bias", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"] }, { "id": "packed_bias_serial", "priority": 0, "when": ["baseContract", "packedShapeOk", "biasOk", "not present.attentionBiasT", "scalarOk"], "derive": { "workgroupSize": "scalarWg" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "bindings": ["query", "bias", "tokenOffset", "cumSeqLen", "output", "params"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] }, { "id": "packed_bias_attn_bias_flash", "priority": 40, "when": ["baseContract", "packedShapeOk", "biasOk", "attnBiasOk", "clusterOk", "clusterSubgroupsOk"], "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"] }, { "id": "packed_bias_attn_bias_flash_nosg", "priority": 20, "when": ["baseContract", "packedShapeOk", "biasOk", "attnBiasOk", "clusterPortableOk"], "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" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"] }, { "id": "packed_bias_attn_bias_serial", "priority": 0, "when": ["baseContract", "packedShapeOk", "biasOk", "attnBiasOk", "scalarOk"], "derive": { "workgroupSize": "scalarWg" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "bindings": ["query", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] }, { "id": "separate_flash", "priority": 40, "when": ["baseContract", "separateShapeOk", "not present.biasT", "not present.attentionBiasT", "clusterOk", "clusterSubgroupsOk"], "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "key4", "value4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"] }, { "id": "separate_flash_nosg", "priority": 20, "when": ["baseContract", "separateShapeOk", "not present.biasT", "not present.attentionBiasT", "clusterPortableOk"], "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" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "key4", "value4", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"] }, { "id": "separate_serial", "priority": 0, "when": ["baseContract", "separateShapeOk", "not present.biasT", "not present.attentionBiasT", "scalarOk"], "derive": { "workgroupSize": "scalarWg" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "bindings": ["query", "key", "value", "tokenOffset", "cumSeqLen", "output", "params"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] }, { "id": "separate_attn_bias_flash", "priority": 40, "when": ["baseContract", "separateShapeOk", "not present.biasT", "attnBiasOk", "clusterOk", "clusterSubgroupsOk"], "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "key4", "value4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"] }, { "id": "separate_attn_bias_flash_nosg", "priority": 20, "when": ["baseContract", "separateShapeOk", "not present.biasT", "attnBiasOk", "clusterPortableOk"], "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" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "key4", "value4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"] }, { "id": "separate_attn_bias_serial", "priority": 0, "when": ["baseContract", "separateShapeOk", "not present.biasT", "attnBiasOk", "scalarOk"], "derive": { "workgroupSize": "scalarWg" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "bindings": ["query", "key", "value", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] }, { "id": "separate_bias_flash", "priority": 40, "when": ["baseContract", "separateShapeOk", "biasOk", "not present.attentionBiasT", "clusterOk", "clusterSubgroupsOk"], "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "key4", "value4", "bias", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"] }, { "id": "separate_bias_flash_nosg", "priority": 20, "when": ["baseContract", "separateShapeOk", "biasOk", "not present.attentionBiasT", "clusterPortableOk"], "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" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "key4", "value4", "bias", "tokenOffset", "cumSeqLen", "output4", "params"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"] }, { "id": "separate_bias_serial", "priority": 0, "when": ["baseContract", "separateShapeOk", "biasOk", "not present.attentionBiasT", "scalarOk"], "derive": { "workgroupSize": "scalarWg" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "bindings": ["query", "key", "value", "bias", "tokenOffset", "cumSeqLen", "output", "params"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] }, { "id": "separate_bias_attn_bias_flash", "priority": 40, "when": ["baseContract", "separateShapeOk", "biasOk", "attnBiasOk", "clusterOk", "clusterSubgroupsOk"], "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "key4", "value4", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)" } } ], "requires": { "features": ["subgroups"] }, "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"] }, { "id": "separate_bias_attn_bias_flash_nosg", "priority": 20, "when": ["baseContract", "separateShapeOk", "biasOk", "attnBiasOk", "clusterPortableOk"], "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" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Cluster", "shader": "attn-packed-varlen-cluster.wgsl.jinja", "bindings": ["query4", "key4", "value4", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"], "dispatch": { "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)" } } ], "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"] }, { "id": "separate_bias_attn_bias_serial", "priority": 0, "when": ["baseContract", "separateShapeOk", "biasOk", "attnBiasOk", "scalarOk"], "derive": { "workgroupSize": "scalarWg" }, "passes": [ { "id": "main", "name": "PackedMultiHeadAttention.Serial", "shader": "attn-packed-varlen-scalar.wgsl.jinja", "bindings": ["query", "key", "value", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"], "dispatch": { "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)", "y": "numHeads", "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)" } } ] } ] }