kv-cache-xpu

Paged KV cache writes for Intel GPUs, loadable through kernels: new key and value rows scatter into a block-paged cache, in float32 or with int8 quantization that halves cache residency. The reference baseline is an index-gather scatter, matched bit-exactly in float32. Companion to paged-attn-xpu, which reads the layout this writes.

Paged attention needs its cache written in the layout the reader expects, and doing that in eager ops costs an index build and a scatter per step, per layer. On an integrated GPU sharing bandwidth with the CPU, that is worth removing, and quantizing the cache to int8 as it lands halves the residency of the thing that grows without bound.

Measured live on an Intel Iris Xe: the paged scatter at 0.021 ms against an eager indexed write's 0.054 (2.6x), bit-exact in float32, with the int8 variant halving cache residency.

Usage

import torch
from kernels import get_kernel

kv = get_kernel("phanerozoic/kv-cache-xpu", version=1, trust_remote_code=True)

key_cache, value_cache = kv.empty_cache(
    num_blocks=256, block_size=16, num_kv_heads=8, head_size=128, device="xpu"
)
kv.reshape_and_cache(key, value, key_cache, value_cache, slot_mapping)

# int8 variant
kv.reshape_and_cache_quant(key, value, kq, vq, ks, vs, slot_mapping)
approx = kv.dequant_cache(kq, ks)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Method

Caches are [num_blocks, block_size, num_kv_heads, head_size]. A slot index s = block * block_size + offset addresses element (s * num_kv_heads + h) * head_size + d, so the scatter is a single flat index computation with no page indirection at write time.

Quantization is per (token, head) absmax to int8, with scales in [num_blocks * block_size, num_kv_heads] indexed by slot. A negative entry in slot_mapping skips that token, which is how padded batches and evicted sequences are handled without a separate compaction pass.

Measured

Against an eager indexed scatter on Intel Iris Xe (Gen12LP, 96 EUs), torch 2.13.0+xpu, 30 iterations after 5 warmup, 8 tokens, 8 KV heads, head_size 128:

The table is reproduced by benchmarks/bench.py, which loads this kernel from the Hub and times it against the eager equivalent on the local device.

Validation

Eight checks on Intel Iris Xe, torch 2.13.0+xpu. Float32 scatter is bit-exact against an index-gather reference for both key and value, and slots outside slot_mapping are verified untouched. Negative slots are confirmed skipped, with the rest of the cache byte-identical to a single expected write. Int8 scales match a reference absmax to 1.86e-09 and the quantized codes are bit-exact against round-to-nearest with clamping; round-trip error is verified at or below half a quantization step for every element.

Requirements and limits

• float32 or float16 inputs, int8 quantized caches. The scatter moves values without computing on them, so float16 is bit-exact; float16 requires aspect::fp16.
• Intel GPU with a working Level Zero driver and a torch XPU build.
• Per (token, head) absmax only; no group-wise scaling within a row.
• The ops ship fake implementations, so they trace under torch.compile without breaking the graph.

References

The paged KV cache layout of vLLM; per-token absmax int8 quantization.

License

Apache-2.0.