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

The deployment of Large Language Models (LLMs) is increasingly shifting toward cost-efficient accelerators such as Google’s Tensor Processing Units (TPUs), where performance per Total Cost of Ownership (TCO) is critical for broad adoption. However, existing kernel research remains largely GPU-centric, leaving a significant gap in high-efficiency TPU inference, particularly for the dynamism and raggedness inherent in modern serving scenarios where traditional XLA-based approaches offer limited support. We introduce Ragged Paged Attention (RPA), a high-performance and flexible attention kernel implemented with Pallas and Mosaic. RPA achieves high hardware utilization by enforcing fine-grained tiling to enable efficient dynamic slicing for irregular memory access and employs a custom software pipeline to fuse KV cache updates directly with attention computation, eliminating a major overhead in the JAX/XLA path. Furthermore, RPA utilizes a distribution-aware compilation strategy to map varying sequence lengths to specialized kernels for decode, fixed-size prefill, and mixed-batch workloads. Evaluated on the Llama 3 8B model on TPU7x, RPA achieves up to 86% memory bandwidth utilization (MBU) in decode and 73% model FLOPs utilization (MFU) in prefill, delivering 60–80× and 16× speedups over JAX/XLA baselines, respectively. RPA has been integrated as the primary TPU backend for leading LLM serving frameworks, including vLLM and SGLang, providing a production-grade foundation for LLM inference on TPU.