arXiv cs.CL· Zishuo Bao, Zeshen Zhang, Linkun Chen, Yucheng Lu·· 10 小时前AI 评分40
FlyWheel：为 TPU 打造的注意力内核
FlyWheel: Attention Kernels Built for the TPU
AI 导读
FlyWheel 是一套面向 TPU 的快速注意力内核，通过重构 softmax 与三角求逆减少向量计算、协调数据布局提升矩阵操作数复用，并重叠独立计算隐藏延迟。
正文 · AI 翻译
摘要：张量处理单元（TPU）上的注意力内核常常无法充分利用其矩阵单元（MXU），原因包括重复的矩阵操作数加载、向量密集型计算，以及导致 MXU 空闲的依赖关系。我们提出 FlyWheel，一套用于 softmax 注意力和线性注意力的快速 TPU 内核。FlyWheel 重新表述 softmax 和三角求逆以减少向量运算，同时协调数据布局以提高矩阵操作数复用。随后它重叠独立计算以隐藏执行延迟。我们将这些技术应用于 TPU v6e 和 v7x 上的 softmax 注意力和 Gated DeltaNet。FlyWheel 相比最先进的 RPA 内核将 softmax 注意力加速最高 $4.14\times$，相比 Splash Attention 加速 $1.80\times$，在 v6e 上达到峰值 BF16 吞吐量的 81%。它还将 Gated DeltaNet 预填充加速最高 $3.61\times$。集成到服务 Qwen3.8-27B 的 vLLM 中，FlyWheel 将预填充吞吐量提升最高 $1.56\times$，同时保持相当的准确率。
| 评论： | 24 页，10 张图，7 张表 |
| 主题： | 计算与语言（cs.CL）；机器学习（cs.LG） |
| 引用为： | arXiv:2611.01426 [cs.CL] |
| （或 arXiv:2611.01426v1 [cs.CL] 用于此版本） | |
| https://doi.org/10.48550/arXiv.2611.01426 由 DataCite 颁发的 arXiv DOI（待注册） |
提交历史
来自：Zishuo Bao [查看邮箱]
[v1]
2026年11月1日星期日 02:07:51 UTC（500 KB）
来源：arXiv cs.CL · arxiv.org