Ragged paged attention
也称为 ragged paged attention、RPA、RPA v3、batched ragged paged attention
先用大白话
ragged paged attention 是 TPU 的注意力 kernel，在一次启动中处理长度各异的一批请求，从分页 KV cache 读取数据。
技术定义
ragged paged attention（RPA）是 TPU 的 Pallas 注意力 kernel，处理通过 block table 寻址的分页 KV cache 中的变长序列，预先计算 page 元数据并在 batch 内流水线化 page 读取。
工程细节
一批在飞行的请求长度参差不齐，它们的 KV page 散布在 HBM 各处。kernel 把序列打包在一起，预先计算 page 元数据，并对 page 读取做三缓冲，在减少填充的同时让 MXU 保持满载。布局很关键：原始 kernel 沿 head 维度打包 key 和 value，每设备单个 FP8 KV head 时每个 128 lane tile 有一半浪费。sequence-on-lane 布局把 token 放到 lane 轴、head 维度放到 sublane 轴，可用 page 翻倍，并允许 head dim 为 64。另一项修复把 KV 计算块与预取块分开，让预取流水线有 VMEM 可以跑在 MXU 前面。移除混合模型的 page 尺寸对齐约束后，batched attention 可以使用 256 token 的 page。
为什么重要
对分页 cache 的注意力计算是 TPU tile 几何与 VMEM 预算冲突最直接的地方。这一个 kernel 中的布局和块大小决策，在硬件不变的情况下让可用 KV 容量、TTFT 和 decode 吞吐量变动了数十个百分点。
如何在 InferenceX 中解读
sequence-on-lane 布局把可用 KV page 从 5,141 提高到 10,283，并发 128 下 8k1k 吞吐量提高 16.5%，中位 TTFT 下降 95%。块大小拆分使 Qwen3-0.6B 的 decode 吞吐量从每秒 64.9k token 提高到 96.3k，256 token page 在并发 512 下为 1k8k 带来约 7% 收益。