Xenova HF Staff

sync 2e7068faf55e

2da2397 verified 6 days ago

6.2 kB

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 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 — kernel metadata (id, digests, provenance)
manifest.json — the op contract (source of truth)
test.json — correctness cases
bench.json — benchmark + tuning cases
attn-flash-decode-splitk-merge.wgsl.jinja
paged-attention.wgsl.jinja
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.

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