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