{ "domain": "com.microsoft", "name": "DecoderMaskedMultiHeadAttention", "sinceVersion": 1, "inputs": { "queryT": { "onnx": "query", "dtype": "T", "rank": 3 }, "keyT": { "onnx": "key", "dtype": "T" }, "valueT": { "onnx": "value", "dtype": "T" }, "maskIndexT": { "onnx": "mask_index", "dtype": "M", "rank": 2, "optional": true }, "attentionBiasT": { "onnx": "attention_bias", "dtype": "T", "rank": 4, "optional": true }, "pastKeyT": { "onnx": "past_key", "dtype": "T", "rank": 4, "optional": true }, "pastValueT": { "onnx": "past_value", "dtype": "T", "rank": 4, "optional": true }, "pastSequenceLengthT": { "onnx": "past_sequence_length", "dtype": "M", "rank": 1, "optional": true } }, "outputs": { "outputT": { "onnx": "output", "dtype": "T", "rank": 3, "shape": "[dim(shapes.queryT, 0), 1, dim(shapes.queryT, 2)]" }, "presentKeyT": { "onnx": "present_key", "dtype": "T", "rank": 4, "optional": true, "shape": "shapes.pastKeyT" }, "presentValueT": { "onnx": "present_value", "dtype": "T", "rank": 4, "optional": true, "shape": "shapes.pastValueT" }, "qkT": { "onnx": "qk", "dtype": "QK", "rank": 4, "optional": true, "shape": "[dim(shapes.queryT, 0), attrs.num_heads, 1, dim(shapes.keyT, 2)]" } }, "attributes": { "num_heads": {}, "past_present_share_buffer": {}, "scale": {}, "mask_filter_value": { "default": -10000 }, "output_qk": {} }, "attributeConstraints": { "num_heads": { "required": true } }, "typeConstraints": { "T": ["float32", "float16"], "QK": ["float32", "float16"], "M": ["int32"] }, "tunables": { "WORKGROUP_SIZE": { "default": 256 }, "COPY_WORKGROUP_SIZE": { "default": 256 }, "COPY_ITEMS_PER_LANE": { "default": 8 }, "COPY_MIN_WORKGROUPS": { "default": 64 }, "MAX_SPLITS": { "default": 8 }, "KEYS_PER_SPLIT": { "default": 256 }, "THREAD_QK_OUTPUT_WORKGROUPS": { "default": 2 } }, "derive": { "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16", "batchSize": "dim(shapes.queryT, 0)", "hiddenSize": "dim(shapes.queryT, 2)", "numHeads": "attrs.num_heads", "headSize": "hiddenSize / max(1, numHeads)", "headDimV4": "headSize / 4", "isCross": "not present.pastKeyT and not present.pastValueT", "crossKvLen": "dim(shapes.keyT, 2) if (isCross and ranks.keyT == 4) else 0", "cacheCapacity": "dim(shapes.pastKeyT, 2) if present.pastKeyT else 0", "physicalKeys": "crossKvLen if isCross else cacheCapacity", "qkRequested": "attrs.output_qk != 0 if has(attrs, \"output_qk\") else false", "floatDtypeOk": "(tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(tensorDtypes.queryT) and tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT and tensorDtypes.outputT == tensorDtypes.queryT", "geometryOk": "ranks.queryT == 3 and dim(shapes.queryT, 1) == 1 and numHeads > 0 and hiddenSize > 0 and hiddenSize % numHeads == 0 and headSize > 0 and headSize % 4 == 0", "outputShapeOk": "ranks.outputT == 3 and dim(shapes.outputT, 0) == batchSize and dim(shapes.outputT, 1) == 1 and dim(shapes.outputT, 2) == hiddenSize", "workgroupOk": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= deviceWorkgroupCap and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and headDimV4 <= deviceWorkgroupCap and tunables.COPY_WORKGROUP_SIZE > 0 and tunables.COPY_WORKGROUP_SIZE <= deviceWorkgroupCap", "workgroupStorageOk": "32 * headDimV4 + 24 * tunables.WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupStorageSize", "numSplits": "max(1, min(tunables.MAX_SPLITS, ceilDiv(physicalKeys, tunables.KEYS_PER_SPLIT)))", "splitScratchBytes": "batchSize * numHeads * numSplits * headSize * 4", "splitStatsBytes": "2 * batchSize * numHeads * numSplits * 4", "scratchFits": "splitScratchBytes <= device.limits.maxStorageBufferBindingSize and splitScratchBytes <= device.limits.maxBufferSize and splitStatsBytes <= device.limits.maxStorageBufferBindingSize and splitStatsBytes <= device.limits.maxBufferSize", "dispatchFits": "batchSize <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and numSplits <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "coreOk": "floatDtypeOk and geometryOk and outputShapeOk and workgroupOk and workgroupStorageOk and scratchFits and dispatchFits and physicalKeys > 0", "crossOk": "coreOk and isCross and ranks.keyT == 4 and ranks.valueT == 4 and dim(shapes.keyT, 0) == batchSize and dim(shapes.keyT, 1) == numHeads and dim(shapes.keyT, 3) == headSize and sameShape(shapes.valueT, shapes.keyT) and not present.attentionBiasT and not present.pastSequenceLengthT and not present.presentKeyT and not present.presentValueT", "cacheDtypeOk": "tensorDtypes.pastKeyT == tensorDtypes.queryT and tensorDtypes.pastValueT == tensorDtypes.queryT and tensorDtypes.presentKeyT == tensorDtypes.queryT and tensorDtypes.presentValueT == tensorDtypes.queryT and tensorDtypes.pastSequenceLengthT == \"int32\"", "selfOk": "coreOk and not isCross and present.pastValueT and present.pastSequenceLengthT and present.presentKeyT and present.presentValueT and attrs.past_present_share_buffer == 1 and ranks.keyT == 3 and ranks.valueT == 3 and dim(shapes.keyT, 0) == batchSize and dim(shapes.keyT, 1) == 1 and dim(shapes.keyT, 2) == hiddenSize and sameShape(shapes.valueT, shapes.keyT) and ranks.pastKeyT == 4 and dim(shapes.pastKeyT, 0) == batchSize and dim(shapes.pastKeyT, 1) == numHeads and dim(shapes.pastKeyT, 3) == headSize and cacheCapacity > 1 and sameShape(shapes.pastValueT, shapes.pastKeyT) and sameShape(shapes.presentKeyT, shapes.pastKeyT) and sameShape(shapes.presentValueT, shapes.pastKeyT) and numel(shapes.pastSequenceLengthT) == 1 and cacheDtypeOk and not present.qkT and not qkRequested", "maskShapeOk": "present.maskIndexT and ranks.maskIndexT == 2 and dim(shapes.maskIndexT, 0) == batchSize and tensorDtypes.maskIndexT == \"int32\" and (dim(shapes.maskIndexT, 1) == crossKvLen if isCross else (dim(shapes.maskIndexT, 1) >= 1 and dim(shapes.maskIndexT, 1) <= cacheCapacity))", "attnBiasShapeOk": "present.attentionBiasT and ranks.attentionBiasT == 4 and tensorDtypes.attentionBiasT == tensorDtypes.queryT and dim(shapes.attentionBiasT, 2) == 1 and (dim(shapes.attentionBiasT, 0) == 1 or dim(shapes.attentionBiasT, 0) == batchSize) and (dim(shapes.attentionBiasT, 1) == 1 or dim(shapes.attentionBiasT, 1) == numHeads) and dim(shapes.attentionBiasT, 3) >= 1 and dim(shapes.attentionBiasT, 3) <= cacheCapacity", "qkShapeOk": "present.qkT and qkRequested and ranks.qkT == 4 and (tensorDtypes.qkT == \"float32\" or tensorDtypes.qkT == \"float16\") and f16Ok(tensorDtypes.qkT) and dim(shapes.qkT, 0) == batchSize and dim(shapes.qkT, 1) == numHeads and dim(shapes.qkT, 2) == 1 and dim(shapes.qkT, 3) == crossKvLen", "crossPlainOk": "crossOk and not present.maskIndexT and not present.qkT and not qkRequested", "crossMaskOk": "crossOk and maskShapeOk and not present.qkT and not qkRequested", "crossQkOk": "crossOk and not present.maskIndexT and qkShapeOk", "crossMaskQkOk": "crossOk and maskShapeOk and qkShapeOk", "selfPlainOk": "selfOk and not present.maskIndexT and not present.attentionBiasT", "selfMaskOk": "selfOk and maskShapeOk and not present.attentionBiasT", "selfBiasOk": "selfOk and not present.maskIndexT and attnBiasShapeOk", "selfMaskBiasOk": "selfOk and maskShapeOk and attnBiasShapeOk", "subgroupSplitKOk": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize >= 8 and device.adapterInfo.subgroupMinSize % 8 == 0 and device.adapterInfo.subgroupMaxSize % 8 == 0 and device.adapterInfo.subgroupMaxSize <= tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE % device.adapterInfo.subgroupMinSize == 0 and tunables.WORKGROUP_SIZE % device.adapterInfo.subgroupMaxSize == 0", "subgroupSplitKPreferred": "subgroupSplitKOk and not narrowSubgroupRange", "attnScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"", "attnVec4": "\"vec4\" if tensorDtypes.queryT == \"float16\" else \"vec4\"", "attnMaskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"", "qkScalarType": "\"f16\" if tensorDtypes.qkT == \"float16\" else \"f32\"", "attnBiasWidth": "dim(shapes.attentionBiasT, 3) if present.attentionBiasT else 0", "attnBiasHeads": "dim(shapes.attentionBiasT, 1) if present.attentionBiasT else 0", "attnBiasBatch": "dim(shapes.attentionBiasT, 0) if present.attentionBiasT else 0", "cacheWordsV4": "numel(shapes.pastKeyT) / 4 if present.pastKeyT else 0", "retainBlockWords": "tunables.COPY_WORKGROUP_SIZE * tunables.COPY_MIN_WORKGROUPS", "retainOccupancyCap": "pick([[cacheWordsV4 >= 8 * retainBlockWords, 8], [cacheWordsV4 >= 4 * retainBlockWords, 4], [cacheWordsV4 >= 2 * retainBlockWords, 2]], 1)", "retainItemsPerLane": "max(1, min(tunables.COPY_ITEMS_PER_LANE, retainOccupancyCap))" }, "bindings": { "query": { "arg": "queryT", "elementType": "$attnVec4" }, "key_cross": { "arg": "keyT", "name": "key", "elementType": "$attnVec4" }, "value_cross": { "arg": "valueT", "name": "value", "elementType": "$attnVec4" }, "key_cache": { "arg": "presentKeyT", "name": "key", "buffer": "read-only-storage", "elementType": "$attnVec4" }, "value_cache": { "arg": "presentValueT", "name": "value", "buffer": "read-only-storage", "elementType": "$attnVec4" }, "key_row": { "arg": "keyT", "name": "key", "elementType": "$attnVec4" }, "value_row": { "arg": "valueT", "name": "value", "elementType": "$attnVec4" }, "past_key": { "arg": "pastKeyT", "elementType": "$attnVec4" }, "past_value": { "arg": "pastValueT", "elementType": "$attnVec4" }, "present_key": { "arg": "presentKeyT", "elementType": "$attnVec4" }, "present_value": { "arg": "presentValueT", "elementType": "$attnVec4" }, "seqlens_k": { "arg": "pastSequenceLengthT", "name": "seqlens_k", "elementType": "i32" }, "past_sequence_length": { "arg": "pastSequenceLengthT", "elementType": "i32" }, "key_pad_mask": { "arg": "maskIndexT", "name": "key_pad_mask", "elementType": "i32" }, "attn_mask": { "arg": "attentionBiasT", "name": "attn_mask", "elementType": "$attnMaskElement" }, "qk": { "arg": "qkT", "elementType": "$qkScalarType" }, "partial_out": { "scratch": "partialOut", "elementType": "vec4" }, "partial_stats": { "scratch": "partialStats", "elementType": "vec2" }, "partial_out_read": { "scratch": "partialOut", "name": "partial_out", "buffer": "read-only-storage", "elementType": "vec4" }, "partial_stats_read": { "scratch": "partialStats", "name": "partial_stats", "buffer": "read-only-storage", "elementType": "vec2" }, "output": { "arg": "outputT", "elementType": "$attnVec4" }, "params_cross": { "name": "params", "struct": [ { "name": "kvSeq", "type": "u32", "value": "crossKvLen" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] }, "params_self": { "name": "params", "struct": [ { "name": "kvSeq", "type": "u32", "value": "cacheCapacity" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] }, "params_self_bias": { "name": "params", "struct": [ { "name": "kvSeq", "type": "u32", "value": "cacheCapacity" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }, { "name": "maskBatchStride", "type": "u32", "value": "0 if attnBiasBatch == 1 else attnBiasHeads * attnBiasWidth" }, { "name": "maskHeadStride", "type": "u32", "value": "0 if attnBiasHeads == 1 else attnBiasWidth" }, { "name": "maskSeqStride", "type": "u32", "value": "0" } ] }, "params_retain": { "name": "params", "struct": [{ "name": "countV4", "type": "u32", "value": "numel(shapes.pastKeyT) / 4" }] }, "params_append": { "name": "params", "struct": [ { "name": "rowCountV4", "type": "u32", "value": "batchSize * hiddenSize / 4" }, { "name": "capacity", "type": "u32", "value": "cacheCapacity" } ] } }, "variants": [ { "id": "cross_direct", "priority": 31, "when": ["crossPlainOk", "subgroupSplitKOk", "numSplits == 1"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitK", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "emitQk": "present.qkT", "useSubgroups": true, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "params_cross", "output"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ], "derive": { "directOutput": true } }, { "id": "cross", "priority": 30, "when": ["crossPlainOk", "subgroupSplitKOk"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitK", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "emitQk": "present.qkT", "useSubgroups": true, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "partial_out", "partial_stats", "params_cross"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.CrossMerge", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "cross_qk_direct", "priority": 31, "when": ["crossQkOk", "subgroupSplitKOk", "numSplits == 1"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQk", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "emitQk": "present.qkT", "qkScalar": "qkScalarType", "useSubgroups": true, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "qk", "params_cross", "output"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ], "derive": { "directOutput": true } }, { "id": "cross_qk", "priority": 30, "when": ["crossQkOk", "subgroupSplitKOk"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQk", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "emitQk": "present.qkT", "qkScalar": "qkScalarType", "useSubgroups": true, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "qk", "partial_out", "partial_stats", "params_cross"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.CrossMerge", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "cross_mask_direct", "priority": 31, "when": ["crossMaskOk", "subgroupSplitKOk", "numSplits == 1"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitK", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "emitQk": "present.qkT", "useSubgroups": true, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "params_cross", "output"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ], "derive": { "directOutput": true } }, { "id": "cross_mask", "priority": 30, "when": ["crossMaskOk", "subgroupSplitKOk"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitK", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "emitQk": "present.qkT", "useSubgroups": true, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "partial_out", "partial_stats", "params_cross"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.CrossMerge", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "cross_mask_qk_direct", "priority": 31, "when": ["crossMaskQkOk", "subgroupSplitKOk", "numSplits == 1"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQk", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "emitQk": "present.qkT", "qkScalar": "qkScalarType", "useSubgroups": true, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "qk", "params_cross", "output"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ], "derive": { "directOutput": true } }, { "id": "cross_mask_qk", "priority": 30, "when": ["crossMaskQkOk", "subgroupSplitKOk"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQk", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "emitQk": "present.qkT", "qkScalar": "qkScalarType", "useSubgroups": true, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "qk", "partial_out", "partial_stats", "params_cross"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.CrossMerge", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "cross_nosg_direct", "priority": 21, "when": ["crossPlainOk", "numSplits == 1"], "requires": {}, "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "emitQk": "present.qkT", "useSubgroups": false, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "params_cross", "output"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ], "derive": { "directOutput": true } }, { "id": "cross_nosg", "priority": 20, "when": ["crossPlainOk"], "requires": {}, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "emitQk": "present.qkT", "useSubgroups": false, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "partial_out", "partial_stats", "params_cross"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.CrossMergeNoSg", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "cross_qk_nosg_direct", "priority": 21, "when": ["crossQkOk", "numSplits == 1"], "requires": {}, "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQkNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "emitQk": "present.qkT", "qkScalar": "qkScalarType", "useSubgroups": false, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "qk", "params_cross", "output"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ], "derive": { "directOutput": true } }, { "id": "cross_qk_nosg", "priority": 20, "when": ["crossQkOk"], "requires": {}, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQkNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "emitQk": "present.qkT", "qkScalar": "qkScalarType", "useSubgroups": false, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "qk", "partial_out", "partial_stats", "params_cross"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.CrossMergeNoSg", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "cross_mask_nosg_direct", "priority": 21, "when": ["crossMaskOk", "numSplits == 1"], "requires": {}, "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "emitQk": "present.qkT", "useSubgroups": false, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "params_cross", "output"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ], "derive": { "directOutput": true } }, { "id": "cross_mask_nosg", "priority": 20, "when": ["crossMaskOk"], "requires": {}, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "emitQk": "present.qkT", "useSubgroups": false, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "partial_out", "partial_stats", "params_cross"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.CrossMergeNoSg", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "cross_mask_qk_nosg_direct", "priority": 21, "when": ["crossMaskQkOk", "numSplits == 1"], "requires": {}, "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQkNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "emitQk": "present.qkT", "qkScalar": "qkScalarType", "useSubgroups": false, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "qk", "params_cross", "output"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ], "derive": { "directOutput": true } }, { "id": "cross_mask_qk_nosg", "priority": 20, "when": ["crossMaskQkOk"], "requires": {}, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQkNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "emitQk": "present.qkT", "qkScalar": "qkScalarType", "useSubgroups": false, "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE" }, "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "qk", "partial_out", "partial_stats", "params_cross"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.CrossMergeNoSg", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "self", "priority": 30, "when": ["selfPlainOk", "subgroupSplitKOk"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "present_retain", "name": "DecoderMaskedMultiHeadAttention.PresentRetain", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": false, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "itemsPerLane": "retainItemsPerLane", "cacheScalar": "attnScalar" }, "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"], "dispatch": { "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "z": 1 }, "viewAlias": [ { "input": "past_key", "output": "present_key" }, { "input": "past_value", "output": "present_value" } ] }, { "id": "present_append", "name": "DecoderMaskedMultiHeadAttention.PresentAppend", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": true, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "headSizeV4": "headDimV4", "numHeads": "numHeads" }, "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"], "dispatch": { "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.SelfSplitK", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "cacheSeqlens": true, "cacheSeqlensIndex": "\"0u\"", "hasKeyPadMask": "present.maskIndexT", "hasMask": "present.attentionBiasT" }, "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "partial_out", "partial_stats", "params_self"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.SelfMerge", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "self_bias", "priority": 30, "when": ["selfBiasOk", "subgroupSplitKOk"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "present_retain", "name": "DecoderMaskedMultiHeadAttention.PresentRetain", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": false, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "itemsPerLane": "retainItemsPerLane", "cacheScalar": "attnScalar" }, "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"], "dispatch": { "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "z": 1 }, "viewAlias": [ { "input": "past_key", "output": "present_key" }, { "input": "past_value", "output": "present_value" } ] }, { "id": "present_append", "name": "DecoderMaskedMultiHeadAttention.PresentAppend", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": true, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "headSizeV4": "headDimV4", "numHeads": "numHeads" }, "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"], "dispatch": { "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.SelfSplitK", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "cacheSeqlens": true, "cacheSeqlensIndex": "\"0u\"", "hasKeyPadMask": "present.maskIndexT", "hasMask": "present.attentionBiasT" }, "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "attn_mask", "partial_out", "partial_stats", "params_self_bias"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.SelfMerge", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "self_mask", "priority": 30, "when": ["selfMaskOk", "subgroupSplitKOk"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "present_retain", "name": "DecoderMaskedMultiHeadAttention.PresentRetain", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": false, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "itemsPerLane": "retainItemsPerLane", "cacheScalar": "attnScalar" }, "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"], "dispatch": { "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "z": 1 }, "viewAlias": [ { "input": "past_key", "output": "present_key" }, { "input": "past_value", "output": "present_value" } ] }, { "id": "present_append", "name": "DecoderMaskedMultiHeadAttention.PresentAppend", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": true, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "headSizeV4": "headDimV4", "numHeads": "numHeads" }, "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"], "dispatch": { "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.SelfSplitK", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "cacheSeqlens": true, "cacheSeqlensIndex": "\"0u\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "hasMask": "present.attentionBiasT" }, "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "key_pad_mask", "partial_out", "partial_stats", "params_self"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.SelfMerge", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "self_mask_bias", "priority": 30, "when": ["selfMaskBiasOk", "subgroupSplitKOk"], "demoteWhen": ["not subgroupSplitKPreferred"], "requires": { "features": ["subgroups"] }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "present_retain", "name": "DecoderMaskedMultiHeadAttention.PresentRetain", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": false, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "itemsPerLane": "retainItemsPerLane", "cacheScalar": "attnScalar" }, "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"], "dispatch": { "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "z": 1 }, "viewAlias": [ { "input": "past_key", "output": "present_key" }, { "input": "past_value", "output": "present_value" } ] }, { "id": "present_append", "name": "DecoderMaskedMultiHeadAttention.PresentAppend", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": true, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "headSizeV4": "headDimV4", "numHeads": "numHeads" }, "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"], "dispatch": { "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.SelfSplitK", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "cacheSeqlens": true, "cacheSeqlensIndex": "\"0u\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "hasMask": "present.attentionBiasT" }, "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "key_pad_mask", "attn_mask", "partial_out", "partial_stats", "params_self_bias"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.SelfMerge", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "self_nosg", "priority": 20, "when": ["selfPlainOk"], "requires": {}, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "present_retain", "name": "DecoderMaskedMultiHeadAttention.PresentRetain", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": false, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "itemsPerLane": "retainItemsPerLane", "cacheScalar": "attnScalar" }, "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"], "dispatch": { "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "z": 1 }, "viewAlias": [ { "input": "past_key", "output": "present_key" }, { "input": "past_value", "output": "present_value" } ] }, { "id": "present_append", "name": "DecoderMaskedMultiHeadAttention.PresentAppend", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": true, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "headSizeV4": "headDimV4", "numHeads": "numHeads" }, "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"], "dispatch": { "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.SelfSplitKNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "cacheSeqlens": true, "cacheSeqlensIndex": "\"0u\"", "hasKeyPadMask": "present.maskIndexT", "hasMask": "present.attentionBiasT", "useSubgroups": false }, "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "partial_out", "partial_stats", "params_self"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.SelfMergeNoSg", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "self_bias_nosg", "priority": 20, "when": ["selfBiasOk"], "requires": {}, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "present_retain", "name": "DecoderMaskedMultiHeadAttention.PresentRetain", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": false, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "itemsPerLane": "retainItemsPerLane", "cacheScalar": "attnScalar" }, "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"], "dispatch": { "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "z": 1 }, "viewAlias": [ { "input": "past_key", "output": "present_key" }, { "input": "past_value", "output": "present_value" } ] }, { "id": "present_append", "name": "DecoderMaskedMultiHeadAttention.PresentAppend", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": true, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "headSizeV4": "headDimV4", "numHeads": "numHeads" }, "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"], "dispatch": { "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.SelfSplitKNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "cacheSeqlens": true, "cacheSeqlensIndex": "\"0u\"", "hasKeyPadMask": "present.maskIndexT", "hasMask": "present.attentionBiasT", "useSubgroups": false }, "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "attn_mask", "partial_out", "partial_stats", "params_self_bias"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.SelfMergeNoSg", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "self_mask_nosg", "priority": 20, "when": ["selfMaskOk"], "requires": {}, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "present_retain", "name": "DecoderMaskedMultiHeadAttention.PresentRetain", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": false, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "itemsPerLane": "retainItemsPerLane", "cacheScalar": "attnScalar" }, "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"], "dispatch": { "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "z": 1 }, "viewAlias": [ { "input": "past_key", "output": "present_key" }, { "input": "past_value", "output": "present_value" } ] }, { "id": "present_append", "name": "DecoderMaskedMultiHeadAttention.PresentAppend", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": true, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "headSizeV4": "headDimV4", "numHeads": "numHeads" }, "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"], "dispatch": { "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.SelfSplitKNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "cacheSeqlens": true, "cacheSeqlensIndex": "\"0u\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "hasMask": "present.attentionBiasT", "useSubgroups": false }, "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "key_pad_mask", "partial_out", "partial_stats", "params_self"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.SelfMergeNoSg", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] }, { "id": "self_mask_bias_nosg", "priority": 20, "when": ["selfMaskBiasOk"], "requires": {}, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" } ], "passes": [ { "id": "present_retain", "name": "DecoderMaskedMultiHeadAttention.PresentRetain", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": false, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "itemsPerLane": "retainItemsPerLane", "cacheScalar": "attnScalar" }, "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"], "dispatch": { "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)", "z": 1 }, "viewAlias": [ { "input": "past_key", "output": "present_key" }, { "input": "past_value", "output": "present_value" } ] }, { "id": "present_append", "name": "DecoderMaskedMultiHeadAttention.PresentAppend", "shader": "dmmha-present.wgsl.jinja", "derive": { "appendRow": true, "workgroupSize": "tunables.COPY_WORKGROUP_SIZE", "headSizeV4": "headDimV4", "numHeads": "numHeads" }, "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"], "dispatch": { "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)", "z": 1 } }, { "id": "split_attention", "name": "DecoderMaskedMultiHeadAttention.SelfSplitKNoSg", "shader": "attn-flash-decode-splitk.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "qNumHeads": "numHeads", "kvNumHeads": "numHeads", "headDim": "headSize", "headDimV4": "headDimV4", "numSplits": "numSplits", "usesF16": "tensorDtypes.queryT == \"float16\"", "cacheSeqlens": true, "cacheSeqlensIndex": "\"0u\"", "hasKeyPadMask": "present.maskIndexT", "maskFilterValue": "attrs.mask_filter_value", "hasMask": "present.attentionBiasT", "useSubgroups": false }, "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "key_pad_mask", "attn_mask", "partial_out", "partial_stats", "params_self_bias"], "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "merge", "name": "DecoderMaskedMultiHeadAttention.SelfMergeNoSg", "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "derive": { "layout": "\"bhsd\"", "splitQueries": false, "scalar": "attnScalar", "qNumHeads": "numHeads", "headDimV4": "headDimV4", "numSplits": "numSplits" }, "bindings": ["partial_out_read", "partial_stats_read", "output"], "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } } ] } ] }