FlyWheel: Attention Kernels Built for the TPU

概要

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.

参考リンク

• Paper ↗

リンク先のライセンス・再現性・データ条件をご確認ください。自動分類や実装の目安は保証ではありません。

実装を始めるプロンプト

参考リンクを確認したうえで、次の内容をコピーして利用できます。

# 論文・記事
FlyWheel: Attention Kernels Built for the TPU

# 概要
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…

# 原典
https://arxiv.org/abs/2611.01426

上記の手法を実装してください。実行前に依存関係とライセンスを確認してください。