---
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 },
});
```