File size: 6,198 Bytes
62958c8 2da2397 62958c8 2da2397 62958c8 2da2397 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | ---
library_name: kernels
license: apache-2.0
tags:
- kernel
- webgpu
- wgsl
---
# com.microsoft.PagedAttention
`com.microsoft` · ONNX Runtime contrib operator · contrib since_version 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.
See the [ONNX Runtime `PagedAttention` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.PagedAttention) for the reference semantics.
## Inputs
| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
| --- | --- | --- | --- | --- | --- | --- | --- |
| `query` | `queryT` | `T` | same as logical dtype | `2` | — | 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. | required |
| `key` | `keyT` | `T` | same as logical dtype | `2` | — | Keys of shape `(num_tokens, kv_num_heads * head_size)`. Absent means `query` carries packed `[Q\|K\|V]`. | optional |
| `value` | `valueT` | `T` | same as logical dtype | `2` | — | Values of shape `(num_tokens, kv_num_heads * head_size)`. Present exactly when `key` is. | optional |
| `key_cache` | `keyCacheT` | `T` | same as logical dtype | `4` | — | Block-based key cache of shape `(num_blocks, block_size, kv_num_heads, head_size)`, updated in place. | required |
| `value_cache` | `valueCacheT` | `T` | same as logical dtype | `4` | — | Block-based value cache with the same shape as `key_cache`, updated in place. | required |
| `cumulative_sequence_length` | `cumulativeSequenceLengthT` | `S` | `int32` | `1` | — | Exclusive prefix sums of the per-sequence token counts, shape `(batch_size + 1)`; sequence `b` owns packed tokens `[cum[b], cum[b+1])`. | required |
| `past_seqlens` | `pastSeqlensT` | `S` | `int32` | `1` | — | Cached history length per sequence, shape `(batch_size)`. | required |
| `block_table` | `blockTableT` | `S` | `int32` | `2` | — | Physical block index per sequence and logical block, shape `(batch_size, max_blocks_per_sequence)`. | required |
| `slot_mapping` | `slotMappingT` | `S` | `int32` | `1` | — | 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. | optional |
## Outputs
| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
| --- | --- | --- | --- | --- | --- | --- |
| `output` | `outputT` | `T` | `2` | derived; see description | Attention output of shape `(num_tokens, num_heads * head_size)`. | required |
| `key_cache` | `keyCacheT` | `T` | `4` | same as `key_cache` | Optional return alias for the updated in-place key cache. The runtime updates both caches together even when only this alias is requested. | optional |
| `value_cache` | `valueCacheT` | `T` | `4` | same as `value_cache` | Optional return alias for the updated in-place value cache. The runtime updates both caches together even when only this alias is requested. | optional |
## Attributes
Attributes and default values (overridable per request):
| Attribute | Default | Description |
| --- | --- | --- |
| `is_causal` | `1` | 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)`. |
## Type constraints
| Variable | Allowed dtypes |
| --- | --- |
| `T` | `float16` |
| `S` | `int32` |
## Device requirements
Every implementation variant requires `shader-f16`; the package has no variant-level fallback without that capability.
## Files
- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
- [`test.json`](build/webgpu/test.json) — correctness cases
- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
- [`attn-flash-decode-splitk-merge.wgsl.jinja`](build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja)
- [`paged-attention.wgsl.jinja`](build/webgpu/paged-attention.wgsl.jinja)
- [`paged-scatter-kv.wgsl.jinja`](build/webgpu/paged-scatter-kv.wgsl.jinja)
## Use with `@huggingface/kernels`
The loader derives every required output's shape and logical dtype from the manifest contract and this call.
It then allocates the result tensors automatically.
The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
Replace each `*Data` placeholder with a typed array containing the corresponding input data.
```js
import { getKernel } from "@huggingface/kernels";
const kernel = await getKernel("webgpu-kernels/com.microsoft.PagedAttention", { version: 1 });
const { keyCacheT, valueCacheT, outputT } = await kernel({
queryT: { data: queryTData, shape: [2, 4] },
keyT: { data: keyTData, shape: [2, 2] },
valueT: { data: valueTData, shape: [2, 2] },
keyCacheT: { data: keyCacheTData, shape: [3, 2, 1, 2] },
valueCacheT: { data: valueCacheTData, shape: [3, 2, 1, 2] },
cumulativeSequenceLengthT: { data: cumulativeSequenceLengthTData, shape: [2] },
pastSeqlensT: { data: pastSeqlensTData, shape: [1] },
blockTableT: { data: blockTableTData, shape: [1, 3] },
}, {
attrs: { num_heads: 2, kv_num_heads: 1 },
});
```
|