| { |
| "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<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\"", |
| "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"] |
| } |
| ] |
| } |
|
|
