arXiv cs.CL· Zishuo Bao, Zeshen Zhang, Linkun Chen, Yucheng Lu·· 11 小时前AI 评分40
FlyWheel：为 TPU 打造的注意力内核
FlyWheel: Attention Kernels Built for the TPU
AI 导读
FlyWheel 是一套面向 TPU 的快速注意力内核，通过重构 softmax 与三角求逆减少向量计算、协调数据布局提升矩阵操作数复用，并重叠独立计算隐藏延迟。
正文 · 原文
Abstract:Attention kernels on Tensor Processing Units (TPUs) often underutilize their matrix units (MXUs) because of repeated matrix operand loads, vector-intensive computation, and dependencies that leave MXUs idle. We present FlyWheel, a set of fast TPU kernels for both softmax and linear attention. FlyWheel reformulates softmax and triangular inversion to reduce vector work, while coordinating data layouts to improve matrix operand reuse. It then overlaps independent computation to hide execution latency. We apply these techniques to both softmax attention and Gated DeltaNet on TPU v6e and v7x. FlyWheel speeds up softmax attention by up to $4.14\times$ over the state-of-the-art RPA kernel and $1.80\times$ over Splash Attention, reaching 81% of peak BF16 throughput on v6e. It also speeds up Gated DeltaNet prefill by up to $3.61\times$. Integrated into vLLM serving Qwen3.8-27B, FlyWheel improves prefill throughput by up to $1.56\times$ while maintaining comparable accuracy.
| Comments: | 24 pages, 10 figures, 7 tables |
| Subjects: | Computation and Language (cs.CL); Machine Learning (cs.LG) |
| Cite as: | arXiv:2611.01426 [cs.CL] |
| (or arXiv:2611.01426v1 [cs.CL] for this version) | |
| https://doi.org/10.48550/arXiv.2611.01426 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Zishuo Bao [view email]
[v1]
Sun, 1 Nov 2026 02:07:51 UTC (500 KB)
来源：arXiv cs.CL · arxiv.org