paged-attn-xpu

Paged KV attention decode for Intel GPUs, loadable through kernels: one query token per sequence attends over a block-paged key/value cache with grouped-query support. The reference baseline is a dense attention that gathers each sequence's pages and runs softmax in float64, matched to 3.58e-07. Companion to kv-cache-xpu, which writes the layout this reads.

The Kernel Hub carries no paged attention for the xpu backend, so serving a model on Intel hardware means either gathering pages into contiguous tensors before every attention call or not paging at all. This supplies the missing piece, and splits the context across work-groups so a single sequence still fills an integrated GPU that a naive one-group-per-head schedule would leave mostly idle.

Measured live on an Intel Iris Xe: 1.49 ms at one sequence, 7.22 at eight, and 12.08 at thirty-two, against eager page-gather attention at 1.71, 40.8, and 78.5, matching a float64 dense reference to 3.58e-07.

Usage

import torch
from kernels import get_kernel

pa = get_kernel("phanerozoic/paged-attn-xpu", version=1, trust_remote_code=True)

out = pa.paged_attention_decode(
    query,          # [num_seqs, num_heads, head_size]   float32, XPU
    key_cache,      # [num_blocks, block_size, num_kv_heads, head_size]
    value_cache,    # same shape as key_cache
    block_tables,   # [num_seqs, max_blocks]  int64
    context_lens,   # [num_seqs]              int64
    scale=None,     # defaults to 1/sqrt(head_size)
)                   # -> [num_seqs, num_heads, head_size]

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. num_heads must be a multiple of num_kv_heads; the ratio is the grouped-query factor.

API

Method

Split-K flash decoding. The context is partitioned across work-groups, then each work-group walks its partition a chunk at a time: one work-item per context position computes a score, the chunk is reduced twice (max, then sum), and the value accumulation reads V with one work-item per head dimension so consecutive work-items touch consecutive addresses. That is three barriers per chunk rather than per position.

Partials combine in a second pass by the standard online-softmax rescale. The rescale is associative, so the result does not depend on where the context was split. The score vector is never materialized, so context length is bounded by the cache rather than by scratch memory. Cache layout is [num_blocks, block_size, num_kv_heads, head_size], matching kv-cache-xpu.

Measured

Against eager PyTorch attention over gathered pages on Intel Iris Xe (Gen12LP, 96 EUs), torch 2.13.0+xpu, 30 iterations after 5 warmup, 32 query heads over 8 KV heads, head_size 128:

Splitting the context across work-groups is what carries the single-sequence case: with one sequence there are only num_heads work-groups of natural parallelism, which leaves most of the device idle. The table is reproduced by benchmarks/bench.py.

Validation

Eleven checks against a dense-attention reference that gathers each sequence's pages and runs softmax in float64, on Intel Iris Xe, torch 2.13.0+xpu. Maximum absolute error 3.58e-07 across all cases. Coverage: multi-head and grouped-query (8 query heads over 2 KV heads); ragged context lengths within a batch; head_size 64, 128, and 256, the last exceeding the work-group width and so exercising the multi-accumulator path; contexts spanning many blocks; and contexts of 600 and 1500 long enough to be split across work-groups, which verifies the combine pass against an unsplit reference. Two invariants are checked exactly: an empty context yields zeros rather than NaN, and attending over a constant value tensor returns that constant, which holds only if the softmax weights sum to one.

Requirements and limits

• float32 or float16. Scores and the online-softmax state are float32 regardless of storage type, and split partials are float32, so splitting never costs accuracy. float16 requires aspect::fp16.
• head_size must not exceed the device's maximum work-group size.
• Intel GPU with a working Level Zero driver and a torch XPU build.
• Decode only, one query token per sequence.
• The ops ship fake implementations, so they trace under torch.compile without breaking the graph.

References

Dao et al., "FlashAttention-2" (2023) and the flash-decoding split; Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention" (vLLM, 2023).

License

Apache-2.0.