TPU上的注意力核不再空转：FlyWheel让softmax注意力最高提速3.49倍，MXU利用率达81%
核心概要
FlyWheel是一组为TPU原生编写的注意力核，同时覆盖softmax注意力与Gated DeltaNet线性注意力。它针对TPU矩阵单元因重复操作数加载、向量运算密集和依赖停顿而闲置的问题，重整数据布局、把三角求逆与卷积改写成矩阵乘法，并显式交错独立计算。在TPU v6e上，softmax注意力相对RPA-v3最高提速3.49倍、相对Splash Attention最高提速1.56倍，达到峰值BF16吞吐的81%。
深度剖析
FlyWheel把TPU上softmax注意力的矩阵单元利用率推到峰值BF16吞吐的81%，此前同类核在因果预填充基准上只用到roofline吞吐的18%至73%。 在TPU v6e的因果预填充基准中，RPA-v3只达到roofline吞吐的18%至53%，Google Tokamax中的Splash Attention达到35%至73%。 在v6e上FlyWheel达到roofline的81%（744 TFLOP/s），Splash Attention为73%、RPA-v3为53%；在v7x上分别为70%（802 TFLOP/s）、59%和43%。
在2K到128K token范围内，FlyWheel的softmax注意力核在MHA下比RPA-v3快3.11至3.49倍，在GQA下快1.50至1.73倍。 同一区间内，FlyWheel比Splash Attention在v6e上快1.10至1.56倍，在v7x上快1.18至1.80倍。 v7x上对RPA-v3的对应倍率为MHA 3.57至4.14倍、GQA 1.57至1.66倍；从8K起FlyWheel的MHA与GQA吞吐相差不超过1.2%，而RPA-v3的MHA吞吐不到其GQA吞吐的一半。
Gated DeltaNet预填充同样受益：在v6e上FlyWheel比GDN-v3快1.84至2.25倍（1K token）、2.45至3.10倍（4K），8K起为2.88至3.61倍。 在Qwen3.8-27B对应的形状上，FlyWheel在16K token时预填充9.9M token/s，GDN-v3为3.0M token/s。 在单个v7x TensorCore上FlyWheel比GDN-v3快1.56至2.62倍；在16K token的单序列上，v6e上的FlyWheel仍比v7x TensorCore上的GDN-v3快2.3至2.6倍。
接入vLLM服务Qwen3.8-27B后，核级提升转化为应用级预填充吞吐提升，且基准精度保持可比。 在v6e上RULER的加速从4K的1.07倍升到128K的1.23倍、512K的1.53倍，BABILong在128K至512K为1.22至1.56倍。 v6e上两次运行在全部评测的RULER任务与长度上均为100%准确率、BABILong分数相同；Video-MME在2048帧上两者均为70.62，1024帧上相差674题中的一题。
启示与展望
评测覆盖TPU v6e单芯片与v7x单TensorCore上的核级基准，以及v6e-8、v6e-16切片和v7x四芯片主机上的vLLM端到端运行，模型为Qwen3.8-27B，基准为RULER、BABILong与Video-MME。
静态锚点依赖每行首个键值块中存在有效键，且需要按实现的操作数转换与归约留出数值余量；分层求逆在BF16操作数与FP32累加下返回的是浮点近似逆，其应用精度由RULER、BABILong与Video-MME上的对照结果给出。