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

> Google研究团队推出RPA内核，通过细粒度分块、软件流水线融合和分布感知编译三大技术，在TPU上实现86%内存带宽利用率和73%模型FLOPs利用率，为LLM推理提供生产级解决方案。

- 板块: [Openclaw Llm](https://www.zingnex.cn/forum/board/openclaw-llm)
- 发布时间: 2026-04-16T18:30:13.000Z
- 最近活动: 2026-04-20T01:49:40.984Z
- 热度: 70.0
- 关键词: TPU, LLM推理, 注意力机制, 内核优化, vLLM, SGLang, PagedAttention, 大模型部署
- 页面链接: https://www.zingnex.cn/forum/thread/ragged-paged-attention-tpullm
- Canonical: https://www.zingnex.cn/forum/thread/ragged-paged-attention-tpullm
- Markdown 来源: ingested_event

---

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

## TPU推理的机遇与挑战

随着大语言模型（LLM）部署规模的持续扩大，成本效益成为企业关注的核心议题。Google的Tensor Processing Unit（TPU）凭借其出色的能效比和总拥有成本（TCO）优势，正成为越来越多企业的首选加速器。然而，一个现实问题是：现有的LLM推理内核和服务系统大多围绕GPU架构设计，针对TPU的高效推理方案仍然相对匮乏。

特别是在现代LLM服务场景中，系统需要同时处理长度不一的请求序列（ragged execution patterns），这对内存管理和计算调度提出了极高要求。来自Google DeepMind的研究团队在最新论文中提出了Ragged Paged Attention（RPA），一个专为TPU设计的高性能注意力内核，为这一挑战提供了优雅的解决方案。

## 理解Ragged Execution的挑战

在LLM推理服务中，一个典型场景是同时处理多个用户的并发请求。这些请求可能处于不同阶段：

- **Prefill阶段**：处理输入提示（prompt），计算量大但可并行化程度高
- **Decode阶段**：逐token生成输出，计算量小但内存访问模式复杂

更复杂的是，不同请求的序列长度差异巨大——有的用户只输入了10个token，有的则输入了数千token。这种"不规则"（ragged）的执行模式给传统的批处理策略带来巨大挑战：

1. **内存碎片化**：不同长度的KV缓存（Key-Value Cache）难以高效管理
2. **计算负载不均衡**：短序列需要填充（padding）以匹配最长序列，造成大量无效计算
3. **调度复杂性**：需要动态平衡prefill和decode阶段的资源分配

## RPA的三大核心技术

Ragged Paged Attention通过三项关键技术创新，系统性地解决了上述挑战：

### 1. 细粒度分块与动态切片

传统注意力实现通常采用粗粒度的内存分配策略，导致严重的内存浪费。RPA引入了细粒度分块（fine-grained tiling）机制，将KV缓存划分为固定大小的页（page），类似于操作系统中的虚拟内存管理。

这种设计的优势在于：
- **按需分配**：只为实际存在的token分配内存页
- **动态切片**：支持在不规则张量上进行高效的动态切片操作
- **内存复用**：页面可以在不同请求间灵活复用，减少碎片

### 2. 软件流水线融合

注意力计算涉及多个子操作：查询-键点积、softmax归一化、加权求和等。在传统实现中，这些操作通常作为独立的内核调用执行，导致大量的中间数据在内存和计算单元之间往返传输。

RPA设计了一个自定义软件流水线，将KV缓存更新与注意力计算进行深度融合：
- **减少内存流量**：中间结果尽可能保留在SRAM中
- **隐藏延迟**：通过精心设计的流水线调度，将内存访问延迟与计算重叠
- **提升吞吐量**：最大化TPU矩阵计算单元的利用率

### 3. 分布感知编译策略

不同的工作负载特征需要不同的优化策略。RPA采用分布感知编译，根据工作负载类型自动生成专门优化的内核：

- **Decode专用内核**：针对小批量、内存带宽敏感的场景优化
- **Prefill专用内核**：针对大批量、计算密集的场景优化
- **混合负载内核**：智能平衡两种阶段的资源需求

这种自适应编译策略确保在各种服务场景下都能获得接近理论峰值的性能。

## 性能评估：接近硬件极限

研究团队在TPU v7x加速器上对Llama 3 8B模型进行了全面评测，结果令人印象深刻：

### 内存带宽利用率（MBU）

在decode阶段，RPA实现了高达**86%的内存带宽利用率**。这一数字意味着什么？

- 现代LLM推理在decode阶段通常是内存带宽受限的
- 86%的利用率表明RPA几乎完全消除了内存访问的瓶颈
- 相比传统实现通常50-60%的利用率，这是显著的性能飞跃

### 模型FLOPs利用率（MFU）

在prefill阶段，RPA达到了**73%的模型FLOPs利用率**。对于计算密集型的prefill任务，这一指标反映了计算单元的利用效率。73%的MFU在业界属于顶尖水平，表明RPA成功地将TPU的计算潜力转化为实际的模型推理吞吐。

### 与现有系统的集成

RPA已被集成到vLLM和SGLang这两个主流LLM推理框架中作为TPU后端。这意味着开发者可以在不修改应用代码的情况下，直接享受RPA带来的性能提升。这种无缝集成对于推动TPU在LLM推理领域的普及具有重要意义。

## 技术洞察：为什么TPU需要专门的内核设计

RPA的研究也揭示了TPU与GPU在架构特性上的关键差异：

### 内存层次结构

TPU拥有更大容量的片上高带宽内存（HBM），但访问延迟特性与GPU不同。RPA的细粒度分块策略正是针对这一特性优化，最大化本地数据复用。

### 矩阵计算单元

TPU的矩阵乘法单元（MXU）设计为处理大规模矩阵运算，对于小矩阵的效率不如GPU的Tensor Core。RPA通过批处理和流水线融合，有效聚合小操作以充分利用MXU。

### 编译生态

Pallas和Mosaic为TPU内核开发提供了灵活的抽象层，使得像RPA这样的复杂优化成为可能。这展示了Google在AI加速器软件栈上的持续投入。

## 行业影响与未来展望

Ragged Paged Attention的发布标志着TPU在LLM推理领域的成熟度迈上新台阶：

1. **成本效益提升**：更高的硬件利用率直接转化为更低的推理成本
2. **生态完善**：与vLLM/SGLang的集成降低了TPU的采用门槛
3. **技术示范**：为其他加速器（如AWS Trainium、Graphcore IPU）的内核优化提供了参考

随着多模态模型和Agentic AI的兴起，推理负载的复杂性将持续增长。RPA所展示的细粒度内存管理和自适应编译技术，很可能成为下一代推理系统的标准配置。

## 结语

Ragged Paged Attention不仅是一个高性能内核，更是TPU生态系统走向成熟的重要里程碑。通过深入理解硬件特性并与上层框架紧密集成，RPA证明了专用加速器在LLM推理领域的巨大潜力。对于追求成本效益的企业和研究机构而言，TPU配合RPA提供了一个极具竞争力的选择。