Ragged Paged Attention on TPU (Video + Blog)

Chinese Video Link: https://youtu.be/cx-tOvSyjN0?si=i0cdc2cOqUSBiBWD
English Video Link: WIP

Paper: https://arxiv.org/abs/2604.15464

Code: https://github.com/vllm-project/tpu-inference/blob/main/tpu_inference/kernels/ragged_paged_attention/v3/kernel.py

1. TPU vs GPU

image.png

https://jax-ml.github.io/scaling-book/tpus/#what-is-a-tpu

https://jax-ml.github.io/scaling-book/gpus/

TPU 7x key components

  • HBM: off-chip high-bandwidth memory
  • VMEM: on-chip vector memory w/ direct load/store to VREG
  • VREG: Vector register of shape 8x128 w/ 32-bit entries
  • VPU: Vector programmable unit. 8x128 SIMD processor in the TensorCore
  • MXU: Matrix unit. A 256x256 systolic array for dense matmul

TPU’s unique characteristics

  • Less flexible for raggedness
    • Systolic array: large. 256x256 32-bit
    • MXU: large matmul. bf16[8,128] @ bf16[128,128] -> f32[8,128]
  • Single-threaded execution

2. Challenges for ragged workloads

Dynamic workloads (e.g., various sequence lengths) need raggedness

https://docs.google.com/spreadsheets/d/1pNFg5yrgHdp3ryCIJHMMNbQEOWPKs06zHCzOS-ijMc0/edit?gid=0#gid=0

Ragged Top

image.png

  • XLA automated optimization for static workloads
  • Coarse-grained tiling: padding and wasted compute
    • Hard to do fine-grained KV
  • Misalignment btw access patterns and layouts? → Communication overhead (scatter, gather)

3. Systolic array

MK * KN → MN

https://docs.google.com/presentation/d/1QyWewJfXfyhQ6f2XAHkZnETv1YHL9gWH8xD8f0WRiD4/edit?slide=id.p#slide=id.p

image.png

4. Tiled layout & packing

Tiling: split a matrix into several smaller sub-matrices

DMA granularity: 512-byte words (or 128 contiguous 32-bit entries) → tiling cannot be too small

Tiling Example

  • Tensor shape = F32(8,256)
  • Tile shape = T(4,128)
    • 4 “chunks” of 128 entries → matches with DMA; matches with TensorCore width
  • logical: row-major order; physical: tiled order

image.png

Packing

  • For narrow data types (bit-width < 32), TPUs employ packing to improve memory efficiency.
  • Multiple elements from adjacent rows of the logical tensor are packed into a single 32-bit “element”
  • Lower and higher bits correspond to different rows

F32(7, 200) with T(4, 128) in logical tiled memory

image.png

BF16(7, 200) with T(4, 128) in logical tiled memory

image.png

RPA supports efficient dynamic slicing and scatter/gather operations

  • placing ragged dimensions along non-tiled axes
  • inserting packing as the second minor dimension

https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#gather

https://gdymind.com/2026/02/07/5D-parallelism-in-LLM-training/

5. Single-threaded, double-buffered

  • Single-thread execution
  • Double buffer: overlap TC compute with DMA

image.png

Pallas kernel

  • DMA, RDMA: for compute-communication overlap
    • Pallas auto-generates multi-buffered pipelines if you don’t explicitly control it
    • RPA does explicitly control for pipelining
  • Register-level (VREG) compute control

6. vLLM Paged Attention

Paged Attention

Paged Attention: partitions each sequence’s KV cache into fixed-size pages

image.png

Challenges on TPU: irregular KV memory accesses

  • KV reads: need to gather data from non-contiguous addresses
  • KV cache updates: ****need to scatter the new KV into partially filled blocks
    • single-token updates during decode

Mixed Batch

Prefill and decode can be mixed

image.png

7. RPA: Ragged Paged Attention

  • Fine-Grained Tiling for Raggedness: a packing dimension makes tiling finer than XLA’s default (7.1–7.4)
  • KV-update-fused Pipeline: overlap KV cache update with attention (7.5–7.7)
  • Distribution-Aware Compilation: different sequence lengths? → different kernels (7.8)
  • Practical notes: tuning (7.9), quantization (7.10)

7.1 Why XLA’s Default Tiling Fails

It’s hard to enforce tiling in XLA.

  • You cannot force it to use T(1, 128) when XLA has its own decision, like T(8, 128).
  • Layouts are fully determined by the XLA’s Layout Assignment pass in HLO
    • Tensor shape + data type → XLA determines tiling and padding

A. Loop dimension shouldn’t be tiled

  • Q shape = [BT, N, H]; K, V shape = [BT, K, H]

image.png

  • The last two dimensions are tiled in XLA. For K, V, they are K, and H in [BT, K, H]
  • However, K is a loop dimension. → It is natural to transpose K and V to [K, BT, H]
    • For KV, this transpose is never materialized. → It is fused into the loading path (see 7.6)

B. Mixed batch: tiling on dynamic dimensions is inefficient

  • BT is dynamic → unpacking, blending, and repacking during dynamic slicing → inefficient tiling

C. Head dimensions: implicit padding issue

  • Head dimensions (number of heads): N in Q and K in K, V

image.png

  • TP and quantization reduce the number of heads in each chip
  • Implicit tiling: cannot reach the minimum T(packing, 128). Example:
    • T(8, 128) tile
    • Tensor shape is BF16(12, 128) → need to pad to BF16(16, 128)
    • However, with a smaller tiling such as T(2, 128), the issue is gone

Solution: RPA introduces a new packing dimension (the 2nd-minor dimension) in Q, K, V → applied to Q in 7.2, to KV in 7.3

7.2 Q Reshape

Since K is the loop dimension (note N = K·G):

  • We express Q as [BT, ⌈KG/P_q⌉, P_q, H]
  • and transpose it to [K, BT, ⌈G/P_q⌉, P_q, H] — grouped by KV head first, so each group’s G query heads are padded to ⌈G/P⌉·P_q separately

Loading Q blocks

  • When loading a Q block to VREGs, we reshape it to [C_q, G, H], allowing multiple query heads within a group to be processed against the same C_k and C_v.
  • As a result, the shared C_k and C_v remain resident in the MXU across all grouped C_q

7.3 KV Merge

  • K and V shapes are the same
  • It’s common that per-chip K becomes 1 after TP — MQA/GQA keep K small to begin with

image.png

  • With BF16, P_k = 2 → implicit padding
  • Solution: merge K and V along the head dimension. Now KV shape is [BT, ⌈2K/P_k⌉, P_k, H]

image.png

Unpacking cost

  • Separating merged C_kv → load 32-bit packed value, then unpack & repack into C_k, C_v
  • Overhead hidden when it doesn’t interfere with MXU scheduling
  • FP8: one 32-bit word holds 2 merged KV elements → repeated extraction per KV head → extra overhead
  • Acceptable trade-off: the smallest update unit (single-token slice) contains K and V for all KV heads → efficient transfers

Other benefits of KV Merge

  • Reduce the # of vector load/store by half
  • Increase the effective DMA transfer size

7.4 Weights: Offline vs Online Processing

Offline idea: reshape & merge projection weights offline → one einsum directly produces Q and merged KV in RPA layouts

  • W_Q: (d_model, N, H) → (K, ⌈G/P⌉, P_q, d_model, H)
  • W_K, W_V: (d_model, K, H) → interleaved into (⌈2K/P_k⌉, P_k, d_model, H)

Why it fails: XLA’s Layout Assignment doesn’t preserve T(packing, 128)

  • Boundary inputs/outputs of a Mosaic kernel are respected, but intermediate tensors get reordered (implicit transposes)
  • Overrides user-specified layout constraints

Solution: defer the transformations to RPA preprocessing (online) to guarantee layout integrity

Future alternatives

  • Projection kernel: constrain layouts in a custom kernel → but ops between projection and RPA may still reorder
  • Kernel fusion: fuse projection into the RPA kernel → layout never altered
  • Layout propagation: if future XLA propagates user layouts to custom kernels → preprocessing eliminated

7.5 Pipeline Overview

image.png

4 types of DMA ops in the pipeline:

  • Fetch B_q: Q block HBM → VMEM, only the effective token range
  • Fetch B_kv: gather KV pages via page table (+ newly projected KV) → contiguous block in VMEM
  • Update U_kv: write newly projected KV back to the KV cache in HBM
  • Send B_o: output block (effective range) VMEM → HBM

Overlapped compute stages:

  1. Load C_q from B_q into VREGs
  2. Strided loads of C_kv from B_kv (transpose fused, see 7.6)
  3. Unpack & repack C_kv → C_k, C_v
  4. Run FlashAttention-2(C_q, C_k, C_v)
  5. Store C_o into VMEM buffer B_o

7.6 KV Transpose Fusion

K is a loop dimension, not part of the matmul → fuse the KV transpose into the loading path via strided vector loads

  • Fetch KV block (b_kv, ⌈2K/P_k⌉, P_k, H) from HBM into VMEM
  • Shapecast to a logical view (b_kv · 2K, H)
  • Strided vector loads of (c_kv, H) with stride 2K → transposed layout on-the-fly, no explicit transpose op

Bank conflicts: arbitrary strides → easily cause VMEM bank conflicts → inefficient load scheduling

  • Fix: add a small offset in the VMEM block layout so the effective stride is odd
  • DMA keeps its effective transfer size; KV loads schedule efficiently

7.7 KV Cache Update Fusion

  • Newly projected KV rides along with Fetch B_kv: appended right after the gathered cache pages, in the same VMEM buffer
  • Update U_kv then reuses that buffer: async write-back to the cache pages, overlapped with attention compute
  • → no separate cache-update kernel; decode’s single-token updates cost no extra pass

7.8 Distribution-Aware Compilation

TensorCore requires a static number of VREGs to unroll a compute block → the unroll must assume the worst-case query length → for decode (q_len = 1), most of it is wasted

Distribution-aware compilation: pass dynamic indices [i, j, k] to the kernel

  • sequences [0, i − 1]: decode-only requests (token length = 1)
  • sequences [i, j − 1]: fixed-size-chunk prefill (where prefill length is static)
  • sequences [j, k − 1]: mixed-batch requests
  • a post-scheduling reordering algorithm guarantees this order: all decode-only and prefill-only requests are placed at the front

Per-case specialization:

  • decode-only & prefill-only: the query block size is known at compile time, allowing the kernel to statically determine VREG usage and schedule computation efficiently
  • mixed: must accommodate unknown effective lengths, potentially leading to wasted compute resources due to padding

Avoiding recompilation

  • JIT enforces static shapes
  • RPA defines the maximum number of tokens and the maximum number of sequences; input data is padded to these upper bounds
  • → shapes stay stable across steps, no recompilation at serving time

7.9 Tuning

4 tunable block sizes:

  • b_q, b_kv: improve DMA–compute overlap
  • c_q, c_kv: reduce VREG pressure / spilling (smaller compute blocks)

Tuning depends on: TPU generation & topology, model (head dim, context length), workload distribution (sequence / generation lengths)

3 sets of block sizes per workload type:

  • Decode-only: memory-bound → saturate memory bandwidth
  • Prefill-only (with fixed-size chunked prefill): compute-bound → maximize MXU utilization, minimize VREG spilling while maintaining high MFU
  • Mixed: hardest; tune for the common case if partially known (e.g., mostly short KV lengths)

Recommendation: tune offline → lookup table → applied at pre-compilation during server bring-up

7.10 Quantization

  • RPA supports per-tensor quantization
  • No per-channel or per-head quantization yet → need more kernel modifications
#JAX #TPU #kernel #Pallas #LLM Inference
Ragged Paged Attention on TPU (Video + Blog)
https://gdymind.github.io/2026/07/06/Ragged-Paged-Attention-on-TPU/
Author
gdymind
Posted on
July 6, 2026
Licensed under