| { |
| "domain": "com.microsoft", |
| "name": "PagedAttention", |
| "sinceVersion": 1, |
| "description": "Attention over a block-based (paged) KV cache: `cumulative_sequence_length` marks the sequence boundaries and `block_table` maps a sequence's history onto scattered blocks. This step's K/V are scattered into the cache, then attended with that history. Grouped-query heads, `scale`, packed `[Q|K|V]`, `slot_mapping`, and float16 cache storage are supported; the cache outputs alias the input caches and are updated in place. Rotary embeddings, softcap, local windows, LATENT layout, narrower value heads, quantized KV, head sinks, q/k normalization, scales, and attention metadata are not implemented.", |
| "inputs": [ |
| { |
| "role": "query", |
| "dtype": "T", |
| "rank": 2, |
| "description": "Packed queries of shape `(num_tokens, num_heads * head_size)`, or `(num_tokens, (num_heads + 2 * kv_num_heads) * head_size)` when `key` and `value` are absent and Q, K and V share one row." |
| }, |
| { |
| "role": "key", |
| "dtype": "T", |
| "rank": 2, |
| "optional": true, |
| "description": "Keys of shape `(num_tokens, kv_num_heads * head_size)`. Absent means `query` carries packed `[Q|K|V]`." |
| }, |
| { |
| "role": "value", |
| "dtype": "T", |
| "rank": 2, |
| "optional": true, |
| "description": "Values of shape `(num_tokens, kv_num_heads * head_size)`. Present exactly when `key` is." |
| }, |
| { |
| "role": "key_cache", |
| "dtype": "T", |
| "rank": 4, |
| "description": "Block-based key cache of shape `(num_blocks, block_size, kv_num_heads, head_size)`, updated in place." |
| }, |
| { |
| "role": "value_cache", |
| "dtype": "T", |
| "rank": 4, |
| "description": "Block-based value cache with the same shape as `key_cache`, updated in place." |
| }, |
| { |
| "role": "cumulative_sequence_length", |
| "dtype": "S", |
| "rank": 1, |
| "description": "Exclusive prefix sums of the per-sequence token counts, shape `(batch_size + 1)`; sequence `b` owns packed tokens `[cum[b], cum[b+1])`." |
| }, |
| { |
| "role": "past_seqlens", |
| "dtype": "S", |
| "rank": 1, |
| "description": "Cached history length per sequence, shape `(batch_size)`." |
| }, |
| { |
| "role": "block_table", |
| "dtype": "S", |
| "rank": 2, |
| "description": "Physical block index per sequence and logical block, shape `(batch_size, max_blocks_per_sequence)`." |
| }, |
| { |
| "role": "slot_mapping", |
| "dtype": "S", |
| "rank": 1, |
| "optional": true, |
| "description": "Flat destination slot, `block_id * block_size + offset`, for each token; `-1` suppresses that token's cache write. When omitted, the slot is derived from `past_seqlens`. `block_table` remains required because it defines the read path." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "output", |
| "dtype": "T", |
| "rank": 2, |
| "shape": "[dim(shapes.queryT, 0), attrs.num_heads * headSize]", |
| "description": "Attention output of shape `(num_tokens, num_heads * head_size)`." |
| }, |
| { |
| "role": "key_cache", |
| "dtype": "T", |
| "rank": 4, |
| "shape": "shapes.keyCacheT", |
| "optional": true, |
| "description": "Optional return alias for the updated in-place key cache. The runtime updates both caches together even when only this alias is requested." |
| }, |
| { |
| "role": "value_cache", |
| "dtype": "T", |
| "rank": 4, |
| "shape": "shapes.valueCacheT", |
| "optional": true, |
| "description": "Optional return alias for the updated in-place value cache. The runtime updates both caches together even when only this alias is requested." |
| } |
| ], |
| "attributes": { "is_causal": 1 }, |
| "attributeDescriptions": { |
| "is_causal": "Whether to apply causal masking. This package supports only value 1. Older ORT schema revisions omit this attribute and are always causal.", |
| "kv_num_heads": "Number of key/value heads.", |
| "num_heads": "Number of query heads.", |
| "scale": "Scale applied to query-key products; zero or omission selects `1 / sqrt(head_size)`." |
| }, |
| "attributeConstraints": { |
| "is_causal": { "values": [1] }, |
| "kv_num_heads": { "required": true }, |
| "num_heads": { "required": true } |
| }, |
| "typeConstraints": { "T": ["float16"], "S": ["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 }, |
| "keyCacheT": { "kind": "tensor", "semantic": "key_cache", "role": "inout" }, |
| "valueCacheT": { "kind": "tensor", "semantic": "value_cache", "role": "inout" }, |
| "cumulativeSequenceLengthT": { |
| "kind": "tensor", |
| "semantic": "cumulative_sequence_length", |
| "role": "input", |
| "dtype": "int32" |
| }, |
| "pastSeqlensT": { "kind": "tensor", "semantic": "past_seqlens", "role": "input", "dtype": "int32" }, |
| "blockTableT": { "kind": "tensor", "semantic": "block_table", "role": "input", "dtype": "int32" }, |
| "slotMappingT": { |
| "kind": "tensor", |
| "semantic": "slot_mapping", |
| "role": "input", |
| "dtype": "int32", |
| "required": false |
| }, |
| "outputT": { "kind": "tensor", "semantic": "output", "role": "output" } |
| }, |
| "tunables": { |
| "WORKGROUP_SIZE": 64, |
| "SCATTER_WORKGROUP_SIZE": 64, |
| "MAX_SPLITS": 16, |
| "SPLIT_MIN_KEYS": 128, |
| "SPLIT_TARGET_WORKGROUPS": 1024 |
| }, |
| "derive": { |
| "tokenCount": "dim(shapes.queryT, 0)", |
| "blockSize": "dim(shapes.keyCacheT, 1)", |
| "headSize": "dim(shapes.keyCacheT, 3)", |
| "maxBlocks": "dim(shapes.blockTableT, 1)", |
| "batchSize": "dim(shapes.cumulativeSequenceLengthT, 0) - 1", |
| "qHidden": "attrs.num_heads * headSize", |
| "kvHidden": "attrs.kv_num_heads * headSize", |
| "qPerKv": "attrs.num_heads / max(1, attrs.kv_num_heads)", |
| "headVec": "headSize / 4", |
| "cacheVec4Ok": "headSize % 4 == 0", |
| "packedStride": "qHidden + 2 * kvHidden", |
| "cacheShapeOk": "ranks.keyCacheT == 4 and ranks.valueCacheT == 4 and sameShape(shapes.valueCacheT, shapes.keyCacheT) and dim(shapes.keyCacheT, 2) == attrs.kv_num_heads and blockSize > 0 and headSize > 0 and tensorDtypes.keyCacheT == tensorDtypes.queryT and tensorDtypes.valueCacheT == tensorDtypes.queryT", |
| "headLayoutOk": "attrs.num_heads > 0 and attrs.kv_num_heads > 0 and attrs.num_heads % attrs.kv_num_heads == 0", |
| "scheduleShapeOk": "ranks.cumulativeSequenceLengthT == 1 and ranks.pastSeqlensT == 1 and ranks.blockTableT == 2 and batchSize >= 1 and dim(shapes.pastSeqlensT, 0) == batchSize and dim(shapes.blockTableT, 0) == batchSize and maxBlocks >= 1", |
| "ioShapeOk": "ranks.queryT == 2 and ranks.outputT == 2 and dim(shapes.outputT, 0) == tokenCount and dim(shapes.outputT, 1) == qHidden and tensorDtypes.outputT == tensorDtypes.queryT and f16Ok(tensorDtypes.queryT)", |
| "separateKv": "present.keyT and present.valueT and ranks.keyT == 2 and ranks.valueT == 2 and dim(shapes.keyT, 0) == tokenCount and dim(shapes.valueT, 0) == tokenCount and dim(shapes.keyT, 1) == kvHidden and dim(shapes.valueT, 1) == kvHidden and tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT and dim(shapes.queryT, 1) == qHidden", |
| "packedKv": "not present.keyT and not present.valueT and dim(shapes.queryT, 1) == packedStride", |
| "slotShapeOk": "dim(shapes.slotMappingT, 0) == tokenCount and ranks.slotMappingT == 1 if present.slotMappingT else true", |
| "pagedContractOk": "cacheShapeOk and headLayoutOk and scheduleShapeOk and ioShapeOk and slotShapeOk", |
| "dispatchFits": "tokenCount <= device.limits.maxComputeWorkgroupsPerDimension and attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension and tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.SCATTER_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup", |
| "pagedBaseWorkgroups": "tokenCount * attrs.kv_num_heads", |
| "pagedKeyCeiling": "maxBlocks * blockSize", |
| "pagedSplitWg": "min(256, max(32, pow2ceil(headVec)))", |
| "pagedSplitCap": "max(1, min(tunables.MAX_SPLITS, pagedKeyCeiling / tunables.SPLIT_MIN_KEYS))", |
| "numSplits": "max(1, min(pagedSplitCap, ceilDiv(tunables.SPLIT_TARGET_WORKGROUPS, max(1, pagedBaseWorkgroups))))", |
| "pagedSplitScratchBytes": "tokenCount * attrs.num_heads * numSplits * headSize * 4", |
| "pagedSplitStatsBytes": "2 * tokenCount * attrs.num_heads * numSplits * 4", |
| "pagedSplitFits": "numSplits >= 2 and cacheVec4Ok and headVec >= 1 and pagedSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and pagedSplitScratchBytes <= device.limits.maxBufferSize and pagedSplitStatsBytes <= device.limits.maxStorageBufferBindingSize and pagedSplitStatsBytes <= device.limits.maxBufferSize and numSplits <= device.limits.maxComputeWorkgroupsPerDimension and headVec <= device.limits.maxComputeInvocationsPerWorkgroup and headVec <= device.limits.maxComputeWorkgroupSizeX and pagedSplitWg <= device.limits.maxComputeInvocationsPerWorkgroup and pagedSplitWg <= device.limits.maxComputeWorkgroupSizeX" |
| }, |
| "constants": { |
| "aScalar": "dtypes.T", |
| "scalar": "dtypes.T", |
| "numHeads": "attrs.num_heads", |
| "kvNumHeads": "attrs.kv_num_heads", |
| "headSize": "headSize", |
| "blockSize": "blockSize", |
| "maxBlocks": "maxBlocks", |
| "qHidden": "qHidden", |
| "kvHidden": "kvHidden", |
| "packedStride": "packedStride", |
| "packedQkv": "not present.keyT", |
| "hasSlotMapping": "present.slotMappingT", |
| "qPerKv": "qPerKv", |
| "attnWorkgroup": "tunables.WORKGROUP_SIZE", |
| "headVec": "headVec", |
| "cacheVec": "(\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\") if cacheVec4Ok else dtypes.T", |
| "cacheVec4Ok": "cacheVec4Ok", |
| "numSplits": "numSplits", |
| "headDimV4": "headVec", |
| "headDim": "headSize", |
| "qNumHeads": "attrs.num_heads", |
| "qHiddenV4": "qHidden / 4", |
| "hasBias": false, |
| "outVec": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\"" |
| }, |
| "bindingSets": { |
| "scatterSeparateDerived": [ |
| { |
| "name": "key", |
| "arg": "keyT", |
| "semantic": "key", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "value", |
| "arg": "valueT", |
| "semantic": "value", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "key_cache", |
| "arg": "keyCacheT", |
| "semantic": "key_cache", |
| "buffer": { "type": "storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "value_cache", |
| "arg": "valueCacheT", |
| "semantic": "value_cache", |
| "buffer": { "type": "storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "cumulative_sequence_length", |
| "arg": "cumulativeSequenceLengthT", |
| "semantic": "cumulative_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "past_seqlens", |
| "arg": "pastSeqlensT", |
| "semantic": "past_seqlens", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "block_table", |
| "arg": "blockTableT", |
| "semantic": "block_table", |
| "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": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" } |
| ] |
| } |
| } |
| ], |
| "scatterSeparateSlot": [ |
| { |
| "name": "key", |
| "arg": "keyT", |
| "semantic": "key", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "value", |
| "arg": "valueT", |
| "semantic": "value", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "key_cache", |
| "arg": "keyCacheT", |
| "semantic": "key_cache", |
| "buffer": { "type": "storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "value_cache", |
| "arg": "valueCacheT", |
| "semantic": "value_cache", |
| "buffer": { "type": "storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "slot_mapping", |
| "arg": "slotMappingT", |
| "semantic": "slot_mapping", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [{ "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }] |
| } |
| } |
| ], |
| "scatterPackedDerived": [ |
| { |
| "name": "query", |
| "arg": "queryT", |
| "semantic": "query", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "key_cache", |
| "arg": "keyCacheT", |
| "semantic": "key_cache", |
| "buffer": { "type": "storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "value_cache", |
| "arg": "valueCacheT", |
| "semantic": "value_cache", |
| "buffer": { "type": "storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "cumulative_sequence_length", |
| "arg": "cumulativeSequenceLengthT", |
| "semantic": "cumulative_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "past_seqlens", |
| "arg": "pastSeqlensT", |
| "semantic": "past_seqlens", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "block_table", |
| "arg": "blockTableT", |
| "semantic": "block_table", |
| "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": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" } |
| ] |
| } |
| } |
| ], |
| "scatterPackedSlot": [ |
| { |
| "name": "query", |
| "arg": "queryT", |
| "semantic": "query", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "key_cache", |
| "arg": "keyCacheT", |
| "semantic": "key_cache", |
| "buffer": { "type": "storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "value_cache", |
| "arg": "valueCacheT", |
| "semantic": "value_cache", |
| "buffer": { "type": "storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "slot_mapping", |
| "arg": "slotMappingT", |
| "semantic": "slot_mapping", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [{ "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }] |
| } |
| } |
| ], |
| "attention": [ |
| { |
| "name": "query", |
| "arg": "queryT", |
| "semantic": "query", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "key_cache", |
| "arg": "keyCacheT", |
| "semantic": "key_cache", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$cacheVec" |
| }, |
| { |
| "name": "value_cache", |
| "arg": "valueCacheT", |
| "semantic": "value_cache", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$cacheVec" |
| }, |
| { |
| "name": "cumulative_sequence_length", |
| "arg": "cumulativeSequenceLengthT", |
| "semantic": "cumulative_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "past_seqlens", |
| "arg": "pastSeqlensT", |
| "semantic": "past_seqlens", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "block_table", |
| "arg": "blockTableT", |
| "semantic": "block_table", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "batchSize" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } |
| ] |
| } |
| } |
| ], |
| "splitAttention": [ |
| { |
| "name": "query", |
| "arg": "queryT", |
| "semantic": "query", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$aScalar" |
| }, |
| { |
| "name": "key_cache", |
| "arg": "keyCacheT", |
| "semantic": "key_cache", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$cacheVec" |
| }, |
| { |
| "name": "value_cache", |
| "arg": "valueCacheT", |
| "semantic": "value_cache", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$cacheVec" |
| }, |
| { |
| "name": "cumulative_sequence_length", |
| "arg": "cumulativeSequenceLengthT", |
| "semantic": "cumulative_sequence_length", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "past_seqlens", |
| "arg": "pastSeqlensT", |
| "semantic": "past_seqlens", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { |
| "name": "block_table", |
| "arg": "blockTableT", |
| "semantic": "block_table", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "i32" |
| }, |
| { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" }, |
| { |
| "name": "partial_stats", |
| "semantic": "partialStats", |
| "buffer": { "type": "storage" }, |
| "elementType": "vec2<f32>" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "batchSize" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } |
| ] |
| } |
| } |
| ], |
| "splitMerge": [ |
| { |
| "name": "partial_out", |
| "semantic": "partialOut", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec4<f32>" |
| }, |
| { |
| "name": "partial_stats", |
| "semantic": "partialStats", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec2<f32>" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outVec" |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "separate_derived_splitk", |
| "description": "Splits each token's key history into contiguous ranges, one workgroup per range, and merges the per-range online-softmax states. It serves grouped-query decode shapes, where one workgroup per `(token, KV head)` leaves too few workgroups to fill the device.", |
| "priority": 10, |
| "requires": { "features": ["shader-f16"] }, |
| "when": ["pagedContractOk", "dispatchFits", "pagedSplitFits", "separateKv", "not present.slotMappingT"], |
| "constants": { "attnWorkgroup": "pagedSplitWg" }, |
| "intermediates": [ |
| { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" }, |
| { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" } |
| ], |
| "passes": [ |
| { |
| "id": "scatter", |
| "name": "PagedAttention.ScatterKV", |
| "shader": "paged-scatter-kv.wgsl.jinja", |
| "bindings": "scatterSeparateDerived", |
| "dispatch": { |
| "threads": "tokenCount * attrs.kv_num_heads * headSize", |
| "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE" |
| } |
| }, |
| { |
| "id": "split_attention", |
| "name": "PagedAttention.AttendSplitK", |
| "source": { "shader": "paged-attention.wgsl.jinja", "inputs": { "splitK": true } }, |
| "bindings": "splitAttention", |
| "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" } |
| }, |
| { |
| "id": "merge", |
| "name": "PagedAttention.AttendSplitKMerge", |
| "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, |
| "bindings": "splitMerge", |
| "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" } |
| } |
| ] |
| }, |
| { |
| "id": "separate_derived", |
| "requires": { "features": ["shader-f16"] }, |
| "when": ["pagedContractOk", "dispatchFits", "separateKv", "not present.slotMappingT"], |
| "passes": [ |
| { |
| "id": "scatter", |
| "name": "PagedAttention.ScatterKV", |
| "shader": "paged-scatter-kv.wgsl.jinja", |
| "bindings": "scatterSeparateDerived", |
| "dispatch": { |
| "threads": "tokenCount * attrs.kv_num_heads * headSize", |
| "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "PagedAttention.Attend", |
| "shader": "paged-attention.wgsl.jinja", |
| "bindings": "attention", |
| "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" } |
| } |
| ] |
| }, |
| { |
| "id": "separate_slot_splitk", |
| "description": "Splits each token's key history into contiguous ranges, one workgroup per range, and merges the per-range online-softmax states. It serves grouped-query decode shapes, where one workgroup per `(token, KV head)` leaves too few workgroups to fill the device.", |
| "priority": 10, |
| "requires": { "features": ["shader-f16"] }, |
| "when": ["pagedContractOk", "dispatchFits", "pagedSplitFits", "separateKv", "present.slotMappingT"], |
| "constants": { "attnWorkgroup": "pagedSplitWg" }, |
| "intermediates": [ |
| { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" }, |
| { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" } |
| ], |
| "passes": [ |
| { |
| "id": "scatter", |
| "name": "PagedAttention.ScatterKV", |
| "shader": "paged-scatter-kv.wgsl.jinja", |
| "bindings": "scatterSeparateSlot", |
| "dispatch": { |
| "threads": "tokenCount * attrs.kv_num_heads * headSize", |
| "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE" |
| } |
| }, |
| { |
| "id": "split_attention", |
| "name": "PagedAttention.AttendSplitK", |
| "source": { "shader": "paged-attention.wgsl.jinja", "inputs": { "splitK": true } }, |
| "bindings": "splitAttention", |
| "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" } |
| }, |
| { |
| "id": "merge", |
| "name": "PagedAttention.AttendSplitKMerge", |
| "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, |
| "bindings": "splitMerge", |
| "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" } |
| } |
| ] |
| }, |
| { |
| "id": "separate_slot", |
| "requires": { "features": ["shader-f16"] }, |
| "when": ["pagedContractOk", "dispatchFits", "separateKv", "present.slotMappingT"], |
| "passes": [ |
| { |
| "id": "scatter", |
| "name": "PagedAttention.ScatterKV", |
| "shader": "paged-scatter-kv.wgsl.jinja", |
| "bindings": "scatterSeparateSlot", |
| "dispatch": { |
| "threads": "tokenCount * attrs.kv_num_heads * headSize", |
| "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "PagedAttention.Attend", |
| "shader": "paged-attention.wgsl.jinja", |
| "bindings": "attention", |
| "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" } |
| } |
| ] |
| }, |
| { |
| "id": "packed_derived_splitk", |
| "description": "Splits each token's key history into contiguous ranges, one workgroup per range, and merges the per-range online-softmax states. It serves grouped-query decode shapes, where one workgroup per `(token, KV head)` leaves too few workgroups to fill the device.", |
| "priority": 10, |
| "requires": { "features": ["shader-f16"] }, |
| "when": ["pagedContractOk", "dispatchFits", "pagedSplitFits", "packedKv", "not present.slotMappingT"], |
| "constants": { "attnWorkgroup": "pagedSplitWg" }, |
| "intermediates": [ |
| { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" }, |
| { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" } |
| ], |
| "passes": [ |
| { |
| "id": "scatter", |
| "name": "PagedAttention.ScatterKV", |
| "shader": "paged-scatter-kv.wgsl.jinja", |
| "bindings": "scatterPackedDerived", |
| "dispatch": { |
| "threads": "tokenCount * attrs.kv_num_heads * headSize", |
| "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE" |
| } |
| }, |
| { |
| "id": "split_attention", |
| "name": "PagedAttention.AttendSplitK", |
| "source": { "shader": "paged-attention.wgsl.jinja", "inputs": { "splitK": true } }, |
| "bindings": "splitAttention", |
| "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" } |
| }, |
| { |
| "id": "merge", |
| "name": "PagedAttention.AttendSplitKMerge", |
| "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, |
| "bindings": "splitMerge", |
| "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" } |
| } |
| ] |
| }, |
| { |
| "id": "packed_derived", |
| "requires": { "features": ["shader-f16"] }, |
| "when": ["pagedContractOk", "dispatchFits", "packedKv", "not present.slotMappingT"], |
| "passes": [ |
| { |
| "id": "scatter", |
| "name": "PagedAttention.ScatterKV", |
| "shader": "paged-scatter-kv.wgsl.jinja", |
| "bindings": "scatterPackedDerived", |
| "dispatch": { |
| "threads": "tokenCount * attrs.kv_num_heads * headSize", |
| "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "PagedAttention.Attend", |
| "shader": "paged-attention.wgsl.jinja", |
| "bindings": "attention", |
| "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" } |
| } |
| ] |
| }, |
| { |
| "id": "packed_slot_splitk", |
| "description": "Splits each token's key history into contiguous ranges, one workgroup per range, and merges the per-range online-softmax states. It serves grouped-query decode shapes, where one workgroup per `(token, KV head)` leaves too few workgroups to fill the device.", |
| "priority": 10, |
| "requires": { "features": ["shader-f16"] }, |
| "when": ["pagedContractOk", "dispatchFits", "pagedSplitFits", "packedKv", "present.slotMappingT"], |
| "constants": { "attnWorkgroup": "pagedSplitWg" }, |
| "intermediates": [ |
| { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" }, |
| { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" } |
| ], |
| "passes": [ |
| { |
| "id": "scatter", |
| "name": "PagedAttention.ScatterKV", |
| "shader": "paged-scatter-kv.wgsl.jinja", |
| "bindings": "scatterPackedSlot", |
| "dispatch": { |
| "threads": "tokenCount * attrs.kv_num_heads * headSize", |
| "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE" |
| } |
| }, |
| { |
| "id": "split_attention", |
| "name": "PagedAttention.AttendSplitK", |
| "source": { "shader": "paged-attention.wgsl.jinja", "inputs": { "splitK": true } }, |
| "bindings": "splitAttention", |
| "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" } |
| }, |
| { |
| "id": "merge", |
| "name": "PagedAttention.AttendSplitKMerge", |
| "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, |
| "bindings": "splitMerge", |
| "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" } |
| } |
| ] |
| }, |
| { |
| "id": "packed_slot", |
| "requires": { "features": ["shader-f16"] }, |
| "when": ["pagedContractOk", "dispatchFits", "packedKv", "present.slotMappingT"], |
| "passes": [ |
| { |
| "id": "scatter", |
| "name": "PagedAttention.ScatterKV", |
| "shader": "paged-scatter-kv.wgsl.jinja", |
| "bindings": "scatterPackedSlot", |
| "dispatch": { |
| "threads": "tokenCount * attrs.kv_num_heads * headSize", |
| "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "PagedAttention.Attend", |
| "shader": "paged-attention.wgsl.jinja", |
| "bindings": "attention", |
| "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
