Ragged paged attention
Also known as RPA, RPA v3, batched ragged paged attention
In plain English
Ragged paged attention is the TPU attention kernel that handles a batch of requests with different lengths reading from a paged KV cache in one launch.
Technical definition
Ragged paged attention is the Pallas attention kernel for TPU that processes variable-length sequences from a block-table-addressed KV cache, precomputing page metadata and pipelining page fetches across a batch.
Engineering details
A batch of in-flight requests has ragged lengths, and their KV pages are scattered across HBM. The kernel batches sequences together, precomputes page metadata, and triple-buffers page fetches to keep the MXU fed while reducing padding. Layout matters: the original kernel packed keys and values along the head dimension, which wasted half of every 128-lane tile for a single FP8 KV head per device. The sequence-on-lane layout puts tokens on the lane axis and the head dimension on sublanes, doubling usable pages and admitting head dimension 64. A separate fix split the KV compute block from the fetch block so the prefetch pipeline had VMEM to run ahead of the MXU. Removing a hybrid page-size alignment constraint let batched attention use a 256-token page.
Why it matters
Attention over the paged cache is where TPU tile geometry and VMEM budget collide most directly. Layout and block-size decisions in this one kernel moved usable KV capacity, TTFT, and decode throughput by tens of percent on unchanged hardware.
How to read it in InferenceX
The sequence-on-lane layout raised usable KV pages from 5,141 to 10,283 and lifted 8k1k throughput 16.5% at concurrency 128 with a 95% cut in median TTFT. The block-size split raised decode throughput from 64.9k to 96.3k tokens per second on Qwen3-0.6B, and the 256-token page gave about 7% on 1k8k at concurrency 512.