Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU

Dimensions overall score 6.0

• Flexible LLM Inference Kernel for TPU Jevin Jiang jevinjiang@google.com Ying Chen chenyin@google.com Blake Hechtman blakehechtman@google.com Fenghui Zhang fhzhang@google.com Yarong Mu ymu@googleclaim 82%
• gap between these "ragged" patterns and XLA’s static-first architecture is a primary goal of this work. At the same time, TPU memory layouts are tiled and coarse-grained, making fine-grained slicing and KV cache updates non-trivialclaim 82%
• compiler Mosaic, exposes low-level primitives for direct control of Direct Memory Access (DMA)/ Remote DMA (RDMA) and register-level computation. This allows developers to rapidly implement custom kernels and workload-specific optimizations, providing a fast time-to-market solutionclaim 82%
• to 86% memory bandwidth utilization (MBU) in decode (memory-bound) and 73% model FLOPs utilization (MFU) in prefill (compute-bound) on TPU7x. Beyond performance, RPA is highly flexible across: (1)Models:supports a wide range of models out of the box, including different sharding strategies, quantization schemesclaim 82%
• multiplication. 128×128 Data 128×128 Data 128×128 Data Zeros 128×128 Data Zeros Contracting Dim Non-contracting Dim 128×128 Data Zeros 128×128 Data Zeros 128×128 Data 128×128 Data 256 256 (a) Contracting size 256 Non-contracting size 256claim 82%
• operations, which enable the overlapping of TensorCore computation with DMAs. This overlap is the primary mechanism for extracting hardware parallelism on TPUs. While Pallas provides high-level abstractions—such as block specifications—to automatically generate multi-buffered pipelinesclaim 82%