Sequence-on-lane 布局
也称为 sequence-on-lane、lane 优先 KV 布局、token-on-lane 布局
先用大白话
Sequence-on-lane 是一种 KV cache 布局,把一页 token 放在 128 宽的向量 lane 轴上而不是 head 维度上,减少浪费在 padding 上的内存。
技术定义
Sequence-on-lane 布局是一种 TPU KV cache 排布方式,把一个 page 的 token 映射到 128 lane 的末维、把 head 维度映射到 sublane 轴,避免小 head 维度放在 lane 轴上时产生的 padding。
FP8、每设备单 KV head 下的可用 KV page
5,141 到 10,283 页;并发 128 的 8k1k 上吞吐量 +16.5%,中位 TTFT -95%
工程细节
TPU 向量单元处理的 tile 末维是 128 lane,末维更小的数组都会被补齐。批量 attention kernel 原先沿 head 维度打包 key 和 value,FP8 下打包因子为 4。每个设备只有一个 KV head 的模型只有两样东西可打包,因此每个 tile 有一半是 padding。把 token 放到 lane 轴上后,无论 head 数量多少 tile 都能填满,head 维度也只需是 32 的倍数而不是 128,head 维度为 64 的模型也能使用该 kernel。
为什么重要
在高并发下,KV 容量就是 decode 吞吐量。这种布局在低并发下每 token 延迟约增加 3%,但在大 batch 时避免了请求排队等待 KV 空间。它还扩大了无需定制工作就能运行在共享 attention kernel 上的模型范围。
如何在 InferenceX 中解读
在报告的 Ironwood 配置中,可用 KV page 从 5,141 翻倍到 10,283。在并发 128 的 8k1k 上,额外容量把吞吐量提高 16.5%,中位 TTFT 降低 95%,因为请求不再等待 KV 空间。