TPU attention kernels stop idling: FlyWheel lifts softmax attention by up to 3.49x and reaches 81% of peak BF16 throughput
Synopsis
FlyWheel is a set of TPU-native attention kernels covering both softmax attention and Gated DeltaNet linear attention. It targets matrix units left idle by repeated operand loads, vector-heavy work and dependency stalls, by coordinating data layouts, recasting triangular inversion and convolution as matrix products, and explicitly interleaving independent computation. On TPU v6e, softmax attention runs up to 3.49x faster than RPA-v3 and up to 1.56x faster than Splash Attention, reaching 81% of peak BF16 throughput.
Interpretation
FlyWheel pushes softmax attention on TPUs to 81% of peak BF16 throughput, where comparable kernels previously used only 18% to 73% of roofline throughput on causal prefill benchmarks. In causal prefill benchmarks on TPU v6e, RPA-v3 reached only 18% to 53% of roofline throughput and Splash Attention in Google's Tokamax reached 35% to 73%. On v6e FlyWheel reaches 81% of roofline (744 TFLOP/s), against 73% for Splash Attention and 53% for RPA-v3; on v7x the figures are 70% (802 TFLOP/s), 59% and 43%.
From 2K to 128K tokens, FlyWheel's softmax attention kernel is 3.11 to 3.49x faster than RPA-v3 under MHA and 1.50 to 1.73x faster under GQA. Over the same range, FlyWheel is 1.10 to 1.56x faster than Splash Attention on v6e and 1.18 to 1.80x on v7x. On v7x the corresponding ratios against RPA-v3 are 3.57 to 4.14x under MHA and 1.57 to 1.66x under GQA; from 8K onward FlyWheel's MHA and GQA throughput differ by at most 1.2%, whereas RPA-v3's MHA throughput is less than half its GQA throughput.
Gated DeltaNet prefill gains as well: on v6e FlyWheel is 1.84 to 2.25x faster than GDN-v3 at 1K tokens, 2.45 to 3.10x at 4K, and 2.88 to 3.61x from 8K on. At the shape matching Qwen3.8-27B, FlyWheel prefills 9.9M tokens/s at 16K tokens, against 3.0M for GDN-v3. On one v7x TensorCore FlyWheel is 1.56 to 2.62x faster than GDN-v3; for a single 16K-token sequence, FlyWheel on a v6e chip is still 2.3 to 2.6x faster than GDN-v3 on a v7x TensorCore.
Integrated into vLLM serving Qwen3.8-27B, the kernel gains carry through to application-level prefill throughput while benchmark accuracy stays comparable. On v6e the RULER speedup rises from 1.07x at 4K to 1.23x at 128K and 1.53x at 512K, and BABILong improves by 1.22 to 1.56x from 128K to 512K. On v6e both runs achieve 100% accuracy on all evaluated RULER tasks and lengths and identical BABILong scores; on Video-MME both score 70.62 at 2048 frames and differ by one of 674 questions at 1024 frames.
Perspective
Evaluation covers kernel benchmarks on one TPU v6e chip and one v7x TensorCore, plus vLLM end-to-end runs on v6e-8 and v6e-16 slices and a four-chip v7x host, with Qwen3.8-27B on RULER, BABILong and Video-MME.
The static anchor relies on a valid key in each row's first key/value tile and needs numerical headroom for the implementation's operand casts and reductions; the hierarchical inverse returns a floating-point approximate inverse under BF16 operands with FP32 accumulation, and its application accuracy is given by the comparisons on RULER, BABILong and Video-MME.