Ragged Paged Attention: 为TPU打造的高性能LLM推理内核

Google研究团队推出Ragged Paged Attention（RPA）内核，专为TPU设计，通过细粒度分块、软件流水线融合、分布感知编译三大技术，实现86%内存带宽利用率和73%模型FLOPs利用率，已集成到vLLM和SGLang框架，为LLM推理提供生产级解决方案，提升TPU在推理场景的成本效益与生态成熟度。

TPU因能效比和总拥有成本（TCO）优势成为企业LLM部署首选，但现有推理方案多针对GPU，TPU高效方案匮乏。现代LLM服务需处理长度不一的请求（ragged执行模式），面临三大挑战：1.内存碎片化（KV缓存管理困难）；2.计算负载不均衡（填充导致无效计算）；3.调度复杂（prefill与decode阶段资源平衡）。

RPA通过三项关键技术解决挑战： 1.细粒度分块与动态切片：将KV缓存划分为固定页，按需分配、动态切片、内存复用，减少碎片； 2.软件流水线融合：深度融合KV更新与注意力计算，中间结果保留在SRAM，隐藏延迟，提升吞吐量； 3.分布感知编译：根据负载类型生成专用内核（decode、prefill、混合负载），自适应优化性能。

在TPU v7x上评测Llama 3 8B模型：

• decode阶段内存带宽利用率达86%（消除内存瓶颈，远超传统50-60%）；
• prefill阶段模型FLOPs利用率达73%（顶尖水平，充分发挥TPU计算潜力）；
• 已集成到vLLM和SGLang作为TPU后端，开发者无需修改代码即可享受性能提升。

RPA针对TPU与GPU的架构差异优化： 1.内存层次：TPU更大HBM，细粒度分块最大化本地数据复用； 2.矩阵计算单元：TPU MXU适合大规模运算，RPA通过批处理和融合聚合小操作； 3.编译生态：Pallas和Mosaic提供灵活抽象，支持复杂内核优化。

RPA标志TPU推理成熟度提升：

• 成本效益：更高硬件利用率降低推理成本；
• 生态完善：与主流框架集成降低TPU采用门槛；
• 技术示范：为其他加速器（如AWS Trainium、Graphcore IPU）提供参考； 未来多模态和Agentic AI负载复杂，RPA的细粒度管理和自适应编译或成下一代推理系统标准。