| { |
| "domain": "com.microsoft", |
| "name": "SparseAttention", |
| "sinceVersion": 1, |
| "description": "Block-sparse causal attention used by Phi-3-small. `block_row_indices` and `block_col_indices` encode one or more CSR block masks, and layouts cycle over query heads. Grouped-query heads, separate or packed `[Q|K|V]`, explicit scaling, partial or full rotary embedding in NeoX or interleaved layout, and float16 are supported. The past/present key and value tensors share allocations and are updated in place. Head sizes must be non-zero multiples of 8; bfloat16 is not implemented.", |
| "inputs": [ |
| { |
| "role": "query", |
| "dtype": "T", |
| "rank": 3, |
| "description": "Query `(batch_size, sequence_length, num_heads * head_size)`, or packed `[Q|K|V]` `(batch_size, sequence_length, (num_heads + 2 * kv_num_heads) * head_size)` when `key` and `value` are omitted." |
| }, |
| { |
| "role": "key", |
| "dtype": "T", |
| "rank": 3, |
| "optional": true, |
| "description": "Key `(batch_size, sequence_length, kv_num_heads * head_size)`. Omitted for packed QKV." |
| }, |
| { |
| "role": "value", |
| "dtype": "T", |
| "rank": 3, |
| "optional": true, |
| "description": "Value `(batch_size, sequence_length, kv_num_heads * head_size)`. Omitted for packed QKV." |
| }, |
| { |
| "role": "past_key", |
| "dtype": "T", |
| "rank": 4, |
| "description": "Key cache `(batch_size, kv_num_heads, max_cache_sequence_length, head_size)`, updated in place." |
| }, |
| { |
| "role": "past_value", |
| "dtype": "T", |
| "rank": 4, |
| "description": "Value cache with the same shape as `past_key`, updated in place." |
| }, |
| { |
| "role": "block_row_indices", |
| "dtype": "M", |
| "rank": 2, |
| "description": "CSR row pointers `(num_layout, max_blocks + 1)`. Each layout starts at zero, is monotonically non-decreasing, and ends no later than that layout's `block_col_indices` width." |
| }, |
| { |
| "role": "block_col_indices", |
| "dtype": "M", |
| "rank": 2, |
| "description": "CSR column indices `(num_layout, max_nnz_blocks)`, right-padded past each layout's non-zero count. Every active entry is in `[0, max_blocks)`." |
| }, |
| { |
| "role": "total_sequence_length", |
| "dtype": "M", |
| "description": "Scalar or one-element vector holding the maximum total key length. Equal to `sequence_length` exactly in the prompt case, which is how the past length is decided. The value fits the cache, the sparse layout's `max_blocks * sparse_block_size` capacity, and the rotary-cache row count when rotary is enabled." |
| }, |
| { |
| "role": "key_total_sequence_lengths", |
| "dtype": "M", |
| "rank": 1, |
| "description": "Per-batch total key length excluding padding, shape `(batch_size)`. Each value is at most `total_sequence_length` and is at least 1 for a prompt or at least `sequence_length` otherwise." |
| }, |
| { |
| "role": "cos_cache", |
| "dtype": "T", |
| "rank": 2, |
| "optional": true, |
| "description": "Rotary cosine cache `(max_rotary_sequence_length, rotary_dimension / 2)`, where the width is a multiple of 8 no larger than `head_size / 2`. Required with `sin_cache` when `do_rotary` is 1." |
| }, |
| { |
| "role": "sin_cache", |
| "dtype": "T", |
| "rank": 2, |
| "optional": true, |
| "description": "Rotary sine cache with the same shape as `cos_cache`; required with it when `do_rotary` is 1." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "output", |
| "dtype": "T", |
| "rank": 3, |
| "shape": "[batchSize, seqLen, numHeads * headSize]", |
| "description": "Attention output `(batch_size, sequence_length, num_heads * head_size)`." |
| }, |
| { |
| "role": "past_key", |
| "dtype": "T", |
| "rank": 4, |
| "shape": "shapes.pastKeyT", |
| "description": "The key cache tensor itself after the in-place append; ONNX names this output `present_key`." |
| }, |
| { |
| "role": "past_value", |
| "dtype": "T", |
| "rank": 4, |
| "shape": "shapes.pastValueT", |
| "description": "The value cache tensor itself after the in-place append; ONNX names this output `present_value`." |
| } |
| ], |
| "attributes": { "do_rotary": 0, "rotary_interleaved": 0 }, |
| "attributeDescriptions": { |
| "num_heads": "Number of query heads.", |
| "kv_num_heads": "Number of key/value heads; must divide `num_heads`.", |
| "sparse_block_size": "Tokens per sparse block; one of 16, 32, 64, 128.", |
| "do_rotary": "Set to 1 to apply rotary embedding to Q and to K before it enters the cache; every other value disables rotary embedding.", |
| "rotary_interleaved": "Set to 1 to rotate adjacent pairs instead of using the NeoX half-split; every other value selects the NeoX layout.", |
| "scale": "Scale applied to query-key products; omitted or zero uses `1 / sqrt(head_size)`." |
| }, |
| "attributeConstraints": { |
| "num_heads": { "required": true }, |
| "kv_num_heads": { "required": true }, |
| "sparse_block_size": { "required": true } |
| }, |
| "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] }, |
| "args": { |
| "queryT": { "kind": "tensor", "semantic": "query", "role": "input" }, |
| "keyT": { "kind": "tensor", "semantic": "key", "role": "input", "required": false }, |
| "valueT": { "kind": "tensor", "semantic": "value", "role": "input", "required": false }, |
| "pastKeyT": { "kind": "tensor", "semantic": "past_key", "role": "inout" }, |
| "pastValueT": { "kind": "tensor", "semantic": "past_value", "role": "inout" }, |
| "blockRowIndicesT": { "kind": "tensor", "semantic": "block_row_indices", "role": "input", "dtype": "int32" }, |
| "blockColIndicesT": { "kind": "tensor", "semantic": "block_col_indices", "role": "input", "dtype": "int32" }, |
| "totalSequenceLengthT": { "kind": "tensor", "semantic": "total_sequence_length", "role": "input", "dtype": "int32" }, |
| "keyTotalSequenceLengthsT": { |
| "kind": "tensor", |
| "semantic": "key_total_sequence_lengths", |
| "role": "input", |
| "dtype": "int32" |
| }, |
| "cosCacheT": { "kind": "tensor", "semantic": "cos_cache", "role": "input", "required": false }, |
| "sinCacheT": { "kind": "tensor", "semantic": "sin_cache", "role": "input", "required": false }, |
| "outputT": { "kind": "tensor", "semantic": "output", "role": "output" } |
| }, |
| "tunables": { |
| "WORKGROUP_SIZE": 128, |
| "APPEND_WORKGROUP_SIZE": 256, |
| "NARROW_MIN_WORKGROUPS": 1024, |
| "QUERY_TILE": 4, |
| "V_STAGE_MAX_WORKGROUPS": 512 |
| }, |
| "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 <= device.limits.maxComputeWorkgroupsPerDimension and batchSize * numHeads <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(ceilDiv(qRotaryElements, tunables.APPEND_WORKGROUP_SIZE), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension 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)", |
| "sparseSgmatLdsBytes": "(64 * 32 + 64 * 64 + 64 * 2 + 128 * 2) * 4", |
| "sparseSgmatGeometryOk": "256 <= device.limits.maxComputeInvocationsPerWorkgroup and 256 <= device.limits.maxComputeWorkgroupSizeX and sgmatQueryTiles <= device.limits.maxComputeWorkgroupsPerDimension and batchSize * numHeads <= device.limits.maxComputeWorkgroupsPerDimension 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" |
| }, |
| "constants": { |
| "scalar": "dtypes.T", |
| "usesF16": "dtypes.T == \"f16\"", |
| "numHeads": "numHeads", |
| "kvNumHeads": "kvNumHeads", |
| "headSize": "headSize", |
| "headVec": "headVec", |
| "cacheVec": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\"", |
| "attnWorkgroup": "sparseAttnWorkgroup", |
| "maxCacheSeq": "maxCacheSeq", |
| "sparseBlockSize": "sparseBlockSize", |
| "numLayout": "numLayout", |
| "maxBlocks": "maxBlocks", |
| "maxNnz": "maxNnz", |
| "packedStride": "packedStride", |
| "packedQkv": "packedQkv", |
| "usesRotary": "useRotary", |
| "rotaryHalf": "rotaryHalf", |
| "rotaryDim": "rotaryDim", |
| "rotaryInterleaved": "rotaryInterleaved", |
| "appendWorkgroupSize": "tunables.APPEND_WORKGROUP_SIZE" |
| }, |
| "bindingSets": { |
| "appendSeparate": [ |
| { |
| "name": "new_key", |
| "arg": "keyT", |
| "semantic": "key", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "new_value", |
| "arg": "valueT", |
| "semantic": "value", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_key", |
| "arg": "pastKeyT", |
| "semantic": "past_key", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_value", |
| "arg": "pastValueT", |
| "semantic": "past_value", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "key_total_sequence_lengths", |
| "arg": "keyTotalSequenceLengthsT", |
| "semantic": "key_total_sequence_lengths", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "total_sequence_length", |
| "arg": "totalSequenceLengthT", |
| "semantic": "total_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "batchSize" }, |
| { "name": "seqLen", "type": "u32", "value": "seqLen" } |
| ] |
| } |
| } |
| ], |
| "appendSeparateRotary": [ |
| { |
| "name": "new_key", |
| "arg": "keyT", |
| "semantic": "key", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "new_value", |
| "arg": "valueT", |
| "semantic": "value", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_key", |
| "arg": "pastKeyT", |
| "semantic": "past_key", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_value", |
| "arg": "pastValueT", |
| "semantic": "past_value", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "key_total_sequence_lengths", |
| "arg": "keyTotalSequenceLengthsT", |
| "semantic": "key_total_sequence_lengths", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "total_sequence_length", |
| "arg": "totalSequenceLengthT", |
| "semantic": "total_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "cos_cache", |
| "arg": "cosCacheT", |
| "semantic": "cos_cache", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "sin_cache", |
| "arg": "sinCacheT", |
| "semantic": "sin_cache", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "batchSize" }, |
| { "name": "seqLen", "type": "u32", "value": "seqLen" } |
| ] |
| } |
| } |
| ], |
| "appendPacked": [ |
| { |
| "name": "packed_qkv", |
| "arg": "queryT", |
| "semantic": "query", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_key", |
| "arg": "pastKeyT", |
| "semantic": "past_key", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_value", |
| "arg": "pastValueT", |
| "semantic": "past_value", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "key_total_sequence_lengths", |
| "arg": "keyTotalSequenceLengthsT", |
| "semantic": "key_total_sequence_lengths", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "total_sequence_length", |
| "arg": "totalSequenceLengthT", |
| "semantic": "total_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "batchSize" }, |
| { "name": "seqLen", "type": "u32", "value": "seqLen" } |
| ] |
| } |
| } |
| ], |
| "appendPackedRotary": [ |
| { |
| "name": "packed_qkv", |
| "arg": "queryT", |
| "semantic": "query", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_key", |
| "arg": "pastKeyT", |
| "semantic": "past_key", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_value", |
| "arg": "pastValueT", |
| "semantic": "past_value", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "key_total_sequence_lengths", |
| "arg": "keyTotalSequenceLengthsT", |
| "semantic": "key_total_sequence_lengths", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "total_sequence_length", |
| "arg": "totalSequenceLengthT", |
| "semantic": "total_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "cos_cache", |
| "arg": "cosCacheT", |
| "semantic": "cos_cache", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "sin_cache", |
| "arg": "sinCacheT", |
| "semantic": "sin_cache", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "batchSize" }, |
| { "name": "seqLen", "type": "u32", "value": "seqLen" } |
| ] |
| } |
| } |
| ], |
| "qRotary": [ |
| { |
| "name": "query", |
| "arg": "queryT", |
| "semantic": "query", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "cos_cache", |
| "arg": "cosCacheT", |
| "semantic": "cos_cache", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "sin_cache", |
| "arg": "sinCacheT", |
| "semantic": "sin_cache", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { "name": "q_rotary", "semantic": "QRotary", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "key_total_sequence_lengths", |
| "arg": "keyTotalSequenceLengthsT", |
| "semantic": "key_total_sequence_lengths", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "total_sequence_length", |
| "arg": "totalSequenceLengthT", |
| "semantic": "total_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "batchSize" }, |
| { "name": "seqLen", "type": "u32", "value": "seqLen" } |
| ] |
| } |
| } |
| ], |
| "attentionDirect": [ |
| { |
| "name": "query", |
| "arg": "queryT", |
| "semantic": "query", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_key", |
| "arg": "pastKeyT", |
| "semantic": "past_key", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$cacheVec" |
| }, |
| { |
| "name": "present_value", |
| "arg": "pastValueT", |
| "semantic": "past_value", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$cacheVec" |
| }, |
| { |
| "name": "block_row_indices", |
| "arg": "blockRowIndicesT", |
| "semantic": "block_row_indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "block_col_indices", |
| "arg": "blockColIndicesT", |
| "semantic": "block_col_indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "key_total_sequence_lengths", |
| "arg": "keyTotalSequenceLengthsT", |
| "semantic": "key_total_sequence_lengths", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "total_sequence_length", |
| "arg": "totalSequenceLengthT", |
| "semantic": "total_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "seqLen", "type": "u32", "value": "seqLen" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } |
| ] |
| } |
| } |
| ], |
| "attentionSgmat": [ |
| { |
| "name": "query", |
| "arg": "queryT", |
| "semantic": "query", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_key", |
| "arg": "pastKeyT", |
| "semantic": "past_key", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_value", |
| "arg": "pastValueT", |
| "semantic": "past_value", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "block_row_indices", |
| "arg": "blockRowIndicesT", |
| "semantic": "block_row_indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "block_col_indices", |
| "arg": "blockColIndicesT", |
| "semantic": "block_col_indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "key_total_sequence_lengths", |
| "arg": "keyTotalSequenceLengthsT", |
| "semantic": "key_total_sequence_lengths", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "total_sequence_length", |
| "arg": "totalSequenceLengthT", |
| "semantic": "total_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "seqLen", "type": "u32", "value": "seqLen" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } |
| ] |
| } |
| } |
| ], |
| "attentionRotary": [ |
| { "name": "q_rotary", "semantic": "QRotary", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "present_key", |
| "arg": "pastKeyT", |
| "semantic": "past_key", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$cacheVec" |
| }, |
| { |
| "name": "present_value", |
| "arg": "pastValueT", |
| "semantic": "past_value", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$cacheVec" |
| }, |
| { |
| "name": "block_row_indices", |
| "arg": "blockRowIndicesT", |
| "semantic": "block_row_indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "block_col_indices", |
| "arg": "blockColIndicesT", |
| "semantic": "block_col_indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "key_total_sequence_lengths", |
| "arg": "keyTotalSequenceLengthsT", |
| "semantic": "key_total_sequence_lengths", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "total_sequence_length", |
| "arg": "totalSequenceLengthT", |
| "semantic": "total_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "seqLen", "type": "u32", "value": "seqLen" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } |
| ] |
| } |
| } |
| ], |
| "attentionSgmatRotary": [ |
| { "name": "q_rotary", "semantic": "QRotary", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "present_key", |
| "arg": "pastKeyT", |
| "semantic": "past_key", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "present_value", |
| "arg": "pastValueT", |
| "semantic": "past_value", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "block_row_indices", |
| "arg": "blockRowIndicesT", |
| "semantic": "block_row_indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "block_col_indices", |
| "arg": "blockColIndicesT", |
| "semantic": "block_col_indices", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "key_total_sequence_lengths", |
| "arg": "keyTotalSequenceLengthsT", |
| "semantic": "key_total_sequence_lengths", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "total_sequence_length", |
| "arg": "totalSequenceLengthT", |
| "semantic": "total_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "seqLen", "type": "u32", "value": "seqLen" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } |
| ] |
| } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "separate", |
| "priority": 0, |
| "when": ["separateContract", "geometryOk"], |
| "passes": [ |
| { |
| "id": "append", |
| "name": "SparseAttention.Append", |
| "shader": "sparse-kv-append.wgsl.jinja", |
| "bindings": "appendSeparate", |
| "constants": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" }, |
| "dispatch": { |
| "threads": "batchSize * kvNumHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "attention", |
| "name": "SparseAttention.Attention", |
| "shader": "sparse-attention.wgsl.jinja", |
| "bindings": "attentionDirect", |
| "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }, |
| "constants": { "qTile": "sparseQueryTile" } |
| } |
| ], |
| "constants": { "vStageWorthIt": "sparseVStageWorthIt" } |
| }, |
| { |
| "id": "separate_sgmat", |
| "priority": 20, |
| "when": ["separateContract", "geometryOk", "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", |
| "bindings": "appendSeparate", |
| "constants": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" }, |
| "dispatch": { |
| "threads": "batchSize * kvNumHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "attention", |
| "name": "SparseAttention.AttentionSgmat", |
| "shader": "sparse-attention-sgmat.wgsl.jinja", |
| "bindings": "attentionSgmat", |
| "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" } |
| } |
| ] |
| }, |
| { |
| "id": "separate_rotary", |
| "priority": 10, |
| "when": ["separateRotaryContract", "geometryOk"], |
| "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }], |
| "passes": [ |
| { |
| "id": "append", |
| "name": "SparseAttention.Append", |
| "shader": "sparse-kv-append.wgsl.jinja", |
| "bindings": "appendSeparateRotary", |
| "constants": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" }, |
| "dispatch": { |
| "threads": "batchSize * kvNumHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "qrotary", |
| "name": "SparseAttention.QueryRotary", |
| "shader": "sparse-q-rotary.wgsl.jinja", |
| "bindings": "qRotary", |
| "dispatch": { |
| "threads": "batchSize * numHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "attention", |
| "name": "SparseAttention.Attention", |
| "shader": "sparse-attention.wgsl.jinja", |
| "bindings": "attentionRotary", |
| "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }, |
| "constants": { "qTile": "sparseQueryTile" } |
| } |
| ], |
| "constants": { "vStageWorthIt": "sparseVStageWorthIt" } |
| }, |
| { |
| "id": "separate_rotary_sgmat", |
| "priority": 30, |
| "when": ["separateRotaryContract", "geometryOk", "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", |
| "bindings": "appendSeparateRotary", |
| "constants": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" }, |
| "dispatch": { |
| "threads": "batchSize * kvNumHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "qrotary", |
| "name": "SparseAttention.QueryRotary", |
| "shader": "sparse-q-rotary.wgsl.jinja", |
| "bindings": "qRotary", |
| "dispatch": { |
| "threads": "batchSize * numHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "attention", |
| "name": "SparseAttention.AttentionSgmat", |
| "shader": "sparse-attention-sgmat.wgsl.jinja", |
| "bindings": "attentionSgmatRotary", |
| "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" } |
| } |
| ] |
| }, |
| { |
| "id": "packed", |
| "priority": 0, |
| "when": ["packedContract", "geometryOk"], |
| "passes": [ |
| { |
| "id": "append", |
| "name": "SparseAttention.Append", |
| "shader": "sparse-kv-append.wgsl.jinja", |
| "bindings": "appendPacked", |
| "constants": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" }, |
| "dispatch": { |
| "threads": "batchSize * kvNumHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "attention", |
| "name": "SparseAttention.Attention", |
| "shader": "sparse-attention.wgsl.jinja", |
| "bindings": "attentionDirect", |
| "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }, |
| "constants": { "qTile": "sparseQueryTile" } |
| } |
| ], |
| "constants": { "vStageWorthIt": "sparseVStageWorthIt" } |
| }, |
| { |
| "id": "packed_sgmat", |
| "priority": 20, |
| "when": ["packedContract", "geometryOk", "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", |
| "bindings": "appendPacked", |
| "constants": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" }, |
| "dispatch": { |
| "threads": "batchSize * kvNumHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "attention", |
| "name": "SparseAttention.AttentionSgmat", |
| "shader": "sparse-attention-sgmat.wgsl.jinja", |
| "bindings": "attentionSgmat", |
| "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" } |
| } |
| ] |
| }, |
| { |
| "id": "packed_rotary", |
| "priority": 10, |
| "when": ["packedRotaryContract", "geometryOk"], |
| "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }], |
| "passes": [ |
| { |
| "id": "append", |
| "name": "SparseAttention.Append", |
| "shader": "sparse-kv-append.wgsl.jinja", |
| "bindings": "appendPackedRotary", |
| "constants": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" }, |
| "dispatch": { |
| "threads": "batchSize * kvNumHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "qrotary", |
| "name": "SparseAttention.QueryRotary", |
| "shader": "sparse-q-rotary.wgsl.jinja", |
| "bindings": "qRotary", |
| "dispatch": { |
| "threads": "batchSize * numHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "attention", |
| "name": "SparseAttention.Attention", |
| "shader": "sparse-attention.wgsl.jinja", |
| "bindings": "attentionRotary", |
| "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }, |
| "constants": { "qTile": "sparseQueryTile" } |
| } |
| ], |
| "constants": { "vStageWorthIt": "sparseVStageWorthIt" } |
| }, |
| { |
| "id": "packed_rotary_sgmat", |
| "priority": 30, |
| "when": ["packedRotaryContract", "geometryOk", "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", |
| "bindings": "appendPackedRotary", |
| "constants": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" }, |
| "dispatch": { |
| "threads": "batchSize * kvNumHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "qrotary", |
| "name": "SparseAttention.QueryRotary", |
| "shader": "sparse-q-rotary.wgsl.jinja", |
| "bindings": "qRotary", |
| "dispatch": { |
| "threads": "batchSize * numHeads * seqLen * headSize", |
| "workgroupSize": "constants.appendWorkgroupSize" |
| } |
| }, |
| { |
| "id": "attention", |
| "name": "SparseAttention.AttentionSgmat", |
| "shader": "sparse-attention-sgmat.wgsl.jinja", |
| "bindings": "attentionSgmatRotary", |
| "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
