{ "domain": "com.microsoft", "name": "SparseAttention", "sinceVersion": 1, "inputs": { "queryT": { "onnx": "query", "dtype": "T", "rank": 3 }, "keyT": { "onnx": "key", "dtype": "T", "rank": 3, "optional": true }, "valueT": { "onnx": "value", "dtype": "T", "rank": 3, "optional": true }, "pastKeyT": { "onnx": "past_key", "dtype": "T", "rank": 4 }, "pastValueT": { "onnx": "past_value", "dtype": "T", "rank": 4 }, "blockRowIndicesT": { "onnx": "block_row_indices", "dtype": "M", "rank": 2, "storage": "int32" }, "blockColIndicesT": { "onnx": "block_col_indices", "dtype": "M", "rank": 2, "storage": "int32" }, "totalSequenceLengthT": { "onnx": "total_sequence_length", "dtype": "M", "storage": "int32" }, "keyTotalSequenceLengthsT": { "onnx": "key_total_sequence_lengths", "dtype": "M", "rank": 1, "storage": "int32" }, "cosCacheT": { "onnx": "cos_cache", "dtype": "T", "rank": 2, "optional": true }, "sinCacheT": { "onnx": "sin_cache", "dtype": "T", "rank": 2, "optional": true } }, "outputs": { "outputT": { "onnx": "output", "dtype": "T", "rank": 3, "shape": "[batchSize, seqLen, numHeads * headSize]" }, "pastKeyT": { "onnx": "past_key", "dtype": "T", "rank": 4, "shape": "shapes.pastKeyT" }, "pastValueT": { "onnx": "past_value", "dtype": "T", "rank": 4, "shape": "shapes.pastValueT" } }, "attributes": { "do_rotary": { "default": 0 }, "rotary_interleaved": { "default": 0 }, "num_heads": {}, "kv_num_heads": {}, "sparse_block_size": {}, "scale": {} }, "attributeConstraints": { "num_heads": { "required": true }, "kv_num_heads": { "required": true }, "sparse_block_size": { "required": true } }, "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] }, "tunables": { "WORKGROUP_SIZE": { "default": 128 }, "APPEND_WORKGROUP_SIZE": { "default": 256 }, "NARROW_MIN_WORKGROUPS": { "default": 1024 }, "QUERY_TILE": { "default": 4 }, "V_STAGE_MAX_WORKGROUPS": { "default": 512 }, "MATRIX_MIN_WORKGROUPS": { "default": 16 } }, "derive": { "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", "batchSize": "dim(shapes.queryT, 0)", "seqLen": "dim(shapes.queryT, 1)", "numHeads": "attrs.num_heads", "kvNumHeads": "attrs.kv_num_heads", "sparseBlockSize": "attrs.sparse_block_size", "headSize": "dim(shapes.pastKeyT, 3)", "headVec": "headSize / 4", "sparseWidthBound": "max(256, tunables.WORKGROUP_SIZE)", "sparseQueryTileCap": "max(1, floor((device.limits.maxComputeWorkgroupStorageSize / 4 - sparseWidthBound) / (2 * headSize + 3 * sparseWidthBound)))", "sparseQueryTileWant": "min(tunables.QUERY_TILE, min(sparseBlockSize, sparseQueryTileCap))", "sparseQueryTile": "1 if seqLen <= 1 else (16 if sparseQueryTileWant >= 16 and seqLen >= 16 else (8 if sparseQueryTileWant >= 8 and seqLen >= 8 else (4 if sparseQueryTileWant >= 4 and seqLen >= 4 else (2 if sparseQueryTileWant >= 2 and seqLen >= 2 else 1))))", "sparseQueryTiles": "ceilDiv(seqLen, sparseQueryTile)", "sparseAttnWorkgroups": "sparseQueryTiles * batchSize * numHeads", "sparseAttnWorkgroup": "min(256, max(32, pow2ceil(headVec))) if sparseAttnWorkgroups >= tunables.NARROW_MIN_WORKGROUPS else tunables.WORKGROUP_SIZE", "maxCacheSeq": "dim(shapes.pastKeyT, 2)", "numLayout": "dim(shapes.blockRowIndicesT, 0)", "maxBlocks": "dim(shapes.blockRowIndicesT, 1) - 1", "maxNnz": "dim(shapes.blockColIndicesT, 1)", "packedQkv": "not present.keyT", "qHidden": "numHeads * headSize", "kvHidden": "kvNumHeads * headSize", "packedStride": "(numHeads + 2 * kvNumHeads) * headSize", "doRotary": "attrs.do_rotary == 1", "rotaryHalf": "dim(shapes.cosCacheT, 1) if doRotary and present.cosCacheT and ranks.cosCacheT == 2 else 0", "rotaryDim": "2 * rotaryHalf", "useRotary": "doRotary and rotaryDim > 0", "rotaryInterleaved": "attrs.rotary_interleaved == 1", "qRotaryElements": "batchSize * numHeads * seqLen * headSize", "cacheShapeOk": "ranks.pastKeyT == 4 and ranks.pastValueT == 4 and dim(shapes.pastKeyT, 0) == batchSize and dim(shapes.pastKeyT, 1) == kvNumHeads and sameShape(shapes.pastValueT, shapes.pastKeyT)", "queryShapeOk": "dim(shapes.queryT, 2) == (packedStride if packedQkv else qHidden)", "kvShapeOk": "packedQkv or (present.valueT and ranks.keyT == 3 and ranks.valueT == 3 and dim(shapes.keyT, 0) == batchSize and dim(shapes.keyT, 1) == seqLen and dim(shapes.keyT, 2) == kvHidden and sameShape(shapes.valueT, shapes.keyT) and tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT)", "kvPairOk": "present.keyT == present.valueT", "rotaryPairOk": "not doRotary or (present.cosCacheT and present.sinCacheT and ranks.cosCacheT == 2 and ranks.sinCacheT == 2 and rotaryHalf % 8 == 0 and rotaryDim <= headSize and sameShape(shapes.sinCacheT, shapes.cosCacheT) and tensorDtypes.cosCacheT == tensorDtypes.queryT and tensorDtypes.sinCacheT == tensorDtypes.queryT)", "blockIndexShapeOk": "ranks.blockRowIndicesT == 2 and ranks.blockColIndicesT == 2 and dim(shapes.blockColIndicesT, 0) == numLayout and maxBlocks >= 1 and maxNnz >= 0 and maxNnz <= maxBlocks * maxBlocks and tensorDtypes.blockRowIndicesT == \"int32\" and tensorDtypes.blockColIndicesT == \"int32\"", "scheduleShapeOk": "(ranks.totalSequenceLengthT == 0 or ranks.totalSequenceLengthT == 1) and numel(shapes.totalSequenceLengthT) == 1 and ranks.keyTotalSequenceLengthsT == 1 and dim(shapes.keyTotalSequenceLengthsT, 0) == batchSize and tensorDtypes.totalSequenceLengthT == \"int32\" and tensorDtypes.keyTotalSequenceLengthsT == \"int32\"", "geometryOk": "tunables.WORKGROUP_SIZE >= 1 and floor(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and tunables.APPEND_WORKGROUP_SIZE >= 1 and floor(tunables.APPEND_WORKGROUP_SIZE) == tunables.APPEND_WORKGROUP_SIZE and tunables.APPEND_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.APPEND_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and sparseQueryTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and batchSize * numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(qRotaryElements, tunables.APPEND_WORKGROUP_SIZE), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and (2 * sparseQueryTile * headSize + (3 * sparseQueryTile + 1) * sparseAttnWorkgroup) * 4 <= device.limits.maxComputeWorkgroupStorageSize and sparseAttnWorkgroup <= device.limits.maxComputeInvocationsPerWorkgroup and sparseAttnWorkgroup <= device.limits.maxComputeWorkgroupSizeX", "contract": "ranks.queryT == 3 and ranks.outputT == 3 and (tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(dtypes.T) and tensorDtypes.pastKeyT == tensorDtypes.queryT and tensorDtypes.pastValueT == tensorDtypes.queryT and tensorDtypes.outputT == tensorDtypes.queryT and numHeads >= 1 and kvNumHeads >= 1 and numHeads % kvNumHeads == 0 and headSize >= 8 and headSize % 8 == 0 and (not doRotary or headSize % 16 == 0) and numLayout >= 1 and numHeads % numLayout == 0 and (sparseBlockSize == 16 or sparseBlockSize == 32 or sparseBlockSize == 64 or sparseBlockSize == 128) and cacheShapeOk and queryShapeOk and kvShapeOk and kvPairOk and rotaryPairOk and blockIndexShapeOk and scheduleShapeOk and dim(shapes.outputT, 0) == batchSize and dim(shapes.outputT, 1) == seqLen and dim(shapes.outputT, 2) == qHidden", "packedContract": "contract and packedQkv and not useRotary", "packedRotaryContract": "contract and packedQkv and useRotary", "separateContract": "contract and not packedQkv and not useRotary", "separateRotaryContract": "contract and not packedQkv and useRotary", "sparseVStageWorthIt": "sparseQueryTiles * batchSize * numHeads <= tunables.V_STAGE_MAX_WORKGROUPS", "sgmatQueryTiles": "ceilDiv(seqLen, 64)", "sgmatDirectQuery": "seqLen % 64 == 0", "sparseSgmatTileN": "64 if (64 * 32 + 64 * 64 + 64 * 2 + 128 * 2) * 4 <= device.limits.maxComputeWorkgroupStorageSize else 32", "sparseSgmatTileK": "sparseSgmatTileN / 2", "sparseSgmatLdsBytes": "(64 * sparseSgmatTileK + 64 * sparseSgmatTileN + 64 * 2 + 128 * 2) * 4", "sparseSgmatGeometryOk": "256 <= device.limits.maxComputeInvocationsPerWorkgroup and 256 <= device.limits.maxComputeWorkgroupSizeX and sgmatQueryTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and batchSize * numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and sparseSgmatLdsBytes <= device.limits.maxComputeWorkgroupStorageSize", "sparseSgmatOk": "tensorDtypes.queryT == \"float32\" and seqLen >= 64 and sparseBlockSize % 64 == 0 and headSize % 32 == 0 and headSize <= 128 and maxCacheSeq % 64 == 0 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and sparseSgmatGeometryOk", "scalar": "dtypes.T", "cacheVec": "\"vec4\" if dtypes.T == \"f16\" else \"vec4\"", "attnWorkgroup": "sparseAttnWorkgroup", "usesRotary": "useRotary", "appendWorkgroupSize": "tunables.APPEND_WORKGROUP_SIZE" }, "when": ["geometryOk"], "bindings": { "new_key": { "arg": "keyT", "buffer": "read-only-storage", "elementType": "$scalar" }, "new_value": { "arg": "valueT", "buffer": "read-only-storage", "elementType": "$scalar" }, "present_key": { "arg": "pastKeyT", "buffer": "storage", "elementType": "$scalar" }, "present_value": { "arg": "pastValueT", "buffer": "storage", "elementType": "$scalar" }, "key_total_sequence_lengths": { "arg": "keyTotalSequenceLengthsT", "buffer": "read-only-storage", "elementType": "i32" }, "total_sequence_length": { "arg": "totalSequenceLengthT", "buffer": "read-only-storage", "elementType": "i32" }, "params": { "buffer": "uniform", "struct": [ { "name": "batchSize", "type": "u32", "value": "batchSize" }, { "name": "seqLen", "type": "u32", "value": "seqLen" } ] }, "cos_cache": { "arg": "cosCacheT", "buffer": "read-only-storage", "elementType": "$scalar" }, "sin_cache": { "arg": "sinCacheT", "buffer": "read-only-storage", "elementType": "$scalar" }, "packed_qkv": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "$scalar" }, "query": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "$scalar" }, "present_key_2": { "arg": "pastKeyT", "name": "present_key", "buffer": "read-only-storage", "elementType": "$cacheVec" }, "present_value_2": { "arg": "pastValueT", "name": "present_value", "buffer": "read-only-storage", "elementType": "$cacheVec" }, "block_row_indices": { "arg": "blockRowIndicesT", "buffer": "read-only-storage", "elementType": "i32" }, "block_col_indices": { "arg": "blockColIndicesT", "buffer": "read-only-storage", "elementType": "i32" }, "output": { "arg": "outputT", "buffer": "storage", "elementType": "$scalar" }, "params_2": { "name": "params", "buffer": "uniform", "struct": [ { "name": "seqLen", "type": "u32", "value": "seqLen" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] }, "q_rotary": { "scratch": "QRotary", "buffer": "read-only-storage", "elementType": "f32" }, "present_key_3": { "arg": "pastKeyT", "name": "present_key", "buffer": "read-only-storage", "elementType": "$scalar" }, "present_value_3": { "arg": "pastValueT", "name": "present_value", "buffer": "read-only-storage", "elementType": "$scalar" }, "q_rotary_2": { "scratch": "QRotary", "name": "q_rotary", "buffer": "storage", "elementType": "f32" } }, "variants": [ { "id": "separate", "priority": 0, "when": ["separateContract"], "derive": { "vStageWorthIt": "sparseVStageWorthIt" }, "passes": [ { "id": "append", "name": "SparseAttention.Append", "shader": "sparse-kv-append.wgsl.jinja", "derive": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" }, "bindings": ["new_key", "new_value", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "attention", "name": "SparseAttention.Attention", "shader": "sparse-attention.wgsl.jinja", "derive": { "qTile": "sparseQueryTile" }, "bindings": ["query", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"], "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" } } ] }, { "id": "separate_sgmat", "priority": 20, "when": ["separateContract", "sparseSgmatOk"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "passes": [ { "id": "append", "name": "SparseAttention.Append", "shader": "sparse-kv-append.wgsl.jinja", "derive": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" }, "bindings": ["new_key", "new_value", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "attention", "name": "SparseAttention.AttentionSgmat", "shader": "sparse-attention-sgmat.wgsl.jinja", "bindings": ["query", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"], "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }, "subgroupCollectivesWidth": 32 } ], "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"] }, { "id": "separate_rotary", "priority": 10, "when": ["separateRotaryContract"], "derive": { "vStageWorthIt": "sparseVStageWorthIt" }, "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }], "passes": [ { "id": "append", "name": "SparseAttention.Append", "shader": "sparse-kv-append.wgsl.jinja", "derive": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" }, "bindings": ["new_key", "new_value", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "cos_cache", "sin_cache", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "qrotary", "name": "SparseAttention.QueryRotary", "shader": "sparse-q-rotary.wgsl.jinja", "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "attention", "name": "SparseAttention.Attention", "shader": "sparse-attention.wgsl.jinja", "derive": { "qTile": "sparseQueryTile" }, "bindings": ["q_rotary", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"], "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" } } ] }, { "id": "separate_rotary_sgmat", "priority": 30, "when": ["separateRotaryContract", "sparseSgmatOk"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }], "passes": [ { "id": "append", "name": "SparseAttention.Append", "shader": "sparse-kv-append.wgsl.jinja", "derive": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" }, "bindings": ["new_key", "new_value", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "cos_cache", "sin_cache", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "qrotary", "name": "SparseAttention.QueryRotary", "shader": "sparse-q-rotary.wgsl.jinja", "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "attention", "name": "SparseAttention.AttentionSgmat", "shader": "sparse-attention-sgmat.wgsl.jinja", "bindings": ["q_rotary", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"], "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }, "subgroupCollectivesWidth": 32 } ], "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"] }, { "id": "packed", "priority": 0, "when": ["packedContract"], "derive": { "vStageWorthIt": "sparseVStageWorthIt" }, "passes": [ { "id": "append", "name": "SparseAttention.Append", "shader": "sparse-kv-append.wgsl.jinja", "derive": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" }, "bindings": ["packed_qkv", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "attention", "name": "SparseAttention.Attention", "shader": "sparse-attention.wgsl.jinja", "derive": { "qTile": "sparseQueryTile" }, "bindings": ["query", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"], "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" } } ] }, { "id": "packed_sgmat", "priority": 20, "when": ["packedContract", "sparseSgmatOk"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "passes": [ { "id": "append", "name": "SparseAttention.Append", "shader": "sparse-kv-append.wgsl.jinja", "derive": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" }, "bindings": ["packed_qkv", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "attention", "name": "SparseAttention.AttentionSgmat", "shader": "sparse-attention-sgmat.wgsl.jinja", "bindings": ["query", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"], "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }, "subgroupCollectivesWidth": 32 } ], "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"] }, { "id": "packed_rotary", "priority": 10, "when": ["packedRotaryContract"], "derive": { "vStageWorthIt": "sparseVStageWorthIt" }, "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }], "passes": [ { "id": "append", "name": "SparseAttention.Append", "shader": "sparse-kv-append.wgsl.jinja", "derive": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" }, "bindings": ["packed_qkv", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "cos_cache", "sin_cache", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "qrotary", "name": "SparseAttention.QueryRotary", "shader": "sparse-q-rotary.wgsl.jinja", "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "attention", "name": "SparseAttention.Attention", "shader": "sparse-attention.wgsl.jinja", "derive": { "qTile": "sparseQueryTile" }, "bindings": ["q_rotary", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"], "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" } } ] }, { "id": "packed_rotary_sgmat", "priority": 30, "when": ["packedRotaryContract", "sparseSgmatOk"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }], "passes": [ { "id": "append", "name": "SparseAttention.Append", "shader": "sparse-kv-append.wgsl.jinja", "derive": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" }, "bindings": ["packed_qkv", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "cos_cache", "sin_cache", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "qrotary", "name": "SparseAttention.QueryRotary", "shader": "sparse-q-rotary.wgsl.jinja", "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"], "dispatch": { "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)", "z": 1 } }, { "id": "attention", "name": "SparseAttention.AttentionSgmat", "shader": "sparse-attention-sgmat.wgsl.jinja", "bindings": ["q_rotary", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"], "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }, "subgroupCollectivesWidth": 32 } ], "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"] } ] }