FlyWheel: Attention Kernels Built for the TPU

Zishuo Bao; Zeshen Zhang; Linkun Chen; Yucheng Lu

cs.CLNov 1, 2026

FlyWheel: Attention Kernels Built for the TPU

Authors: Zishuo Bao, Zeshen Zhang, Linkun Chen, Yucheng Lu

Organizations: New York University · New York University Shanghai · Google

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.

Figures & tables

Figure 1: Per-head tiling with ragged boundaries. (a) Under MHA, RPA-v3 pads each query head and holds 8 tokens per BF16 register; FlyWheel holds 16 tokens of one head without padding. (b) A sequence starting at row 13 is written through a window aligned at row 8; FlyWheel first restores rows 8–12 from the previous VMEM buffer.

Figure 2: Filling dependency-induced execution gaps. (a) RPA-v3 leaves the MXU idle during exponentiation; FlyWheel overlaps each fragment’s scores and exponentials with the previous fragment’s PV product. (b) Gaps between the dependent inverse products 1–8 of one GDN chunk are filled with the other chunk’s stacked product (A) and state-dependent work. Lengths are schematic.

Figure 3: Replacing vector operations with matrix products. (a) Hierarchical triangular inversion merges

4\times 4

blocks into

16\times 16

and then

64\times 64

blocks. (b) A four-tap causal convolution (

C=8

) as one product

\bm{Y}=\bm{S}\bm{W}

, where

\bm{W}

stacks the tap-scaled inputs and halo rows, and the constant selector

\bm{S}

(

\bm{S}_{\mathrm{conv}}

in the text) sums the required rows.

	BF16 peak	VMEM	HBM	HBM BW
	(TFLOP/s)	(MiB)	(GiB)	(TB/s)
v6e chip	918	128	32	1.6
v7x TensorCore	1153	64	96	3.7

Table 1: Hardware of the kernel benchmarks. A v7x chip has two TensorCores; the row lists one of them.

Figure 4: Kernel-only throughput of softmax attention, causal prefill with head dimension 256,

B=1

, and 32 query heads. Rows are one v6e chip and one v7x TensorCore; columns are MHA (32:32) and GQA (32:4). The numbers under each length are the speedups of FlyWheel over each baseline.

Figure 5: The throughput of Figure 4 on the roofline of one TPU v6e chip and one v7x TensorCore, for 1K to 128K tokens in powers of two. Solid lines with circles are MHA (32:32) and dashed lines with triangles are GQA (32:4). The number at the end of each line is its throughput at 128K tokens.

Figure 6: Throughput of the fused causal conv1d and GDN prefill call, with 16 query/key heads and

d_{k}=d_{v}=128

. Numbers under each length are FlyWheel’s speedups over GDN-v3 on the same device. On v7x at

n_{v}=48

, 64-row tiles fit in VMEM but can accommodate only a single chunk, preventing cross-chunk pipelining and resulting in a slight performance drop.

Figure 7: Qwen3.8-27B in vLLM on v6e (top) and v7x (bottom). Numbers under each group are FlyWheel’s speedups over the baseline. The v7x runs use 20 samples per RULER and BABILong task and 8 Video-MME questions.

Configuration	Time (ms)	TFLOP/s
Static anchor, pipeline (FlyWheel)	24.43	720
Online softmax, pipeline	30.21	582
FA4 rescale, pipeline	30.86	570
Static anchor, no pipeline	33.79	521
Online softmax, no pipeline	49.79	353
with the LP LLO scheduler	54.96	320

Table 2: Softmax attention ablations on v6e: causal prefill at 32K tokens,

D=256

, and 32 query and KV heads. Time is kernel device time. All configurations use the same data layout; the ablations vary softmax computation and fragment scheduling.

	$n_{v}$
Configuration / added design	16	32	48
GDN-v3 time (ms)	1.31	2.15	2.75
Coordinated layouts	1.52	1.68	1.59
Matrix operand reuse	1.06	1.22	1.16
Reduced vector work	1.58	1.51	1.56
Explicit pipelining	1.25	1.14	1.13

Table 3: Incremental GDN ablations on v6e: fused causal conv1d and GDN prefill of 8K tokens in one sequence, with 16 query/key heads and

d_{k}=d_{v}=128

. Each design row reports speedup over the preceding configuration; their product gives the total speedup, up to rounding. All designs are described in § 3.2 .

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Softmax attention block time on one v6e chip, from hidden states to hidden states (

d_{\mathrm{model}}=8192

), causal prefill with head dimension 256,

B=1

, and 32 query heads, split by segment and normalized to the slowest block of each group. Panels are MHA (32:32) and GQA (32:4). R, S, and F are RPA-v3 , Splash-Attention , and FlyWheel; the numbers under each length are the wall-clock speedups of FlyWheel’s block over each baseline.

Figure 9: FlyWheel’s kernel throughput over the faster of RPA-v3 and Splash-Attention on one v6e chip, at every cell of the grid: 22 combinations of query heads $H$ , KV heads $H_{kv}$ , and head dimension $D$ , each swept over 8 batch sizes and 8 sequence lengths with $BT\leq 128$ K. Blue is FlyWheel ahead, red behind. Columns share a KV head count, so the $H_{kv}=32$ slots are empty for $H=16$ ; the pale corner of each panel is the part of the rectangle the token budget excludes.

Figure 10: Kernel utilization of the v6e BF16 peak over the same grid. Each line is the median across all KV head counts and batch sizes at that length, and the band is their full range. RPA-v3 ’s band is wide because its throughput depends on the KV head count (§ 4.2 ); FlyWheel’s and Splash-Attention ’s are narrow.

$H_{q}{:}H_{kv}$	Sequences	RPA-v3	FlyWheel	Speedup
32:4	256	5.148	5.023	1.02 $\times$
32:4	128	2.574	2.519	1.02 $\times$
32:4	32	0.660	0.649	1.02 $\times$
32:32	128	22.426	19.875	1.13 $\times$
32:32	32	5.645	4.950	1.14 $\times$

Appendix

Table 4: BF16 decode against RPA-v3 on one v7x TensorCore: 32 query heads,

D=256

, 128-token pages, 15K cached tokens per sequence followed by 1024 decode steps. Time is ms per step.

			v6e chip			v7x TensorCore
Cached	Batch	$T$	FlyWheel	Batched-RPA	Speedup	FlyWheel	Batched-RPA	Speedup
2048	32	1	0.245	0.249	1.02 $\times$	0.103	0.151	1.47 $\times$
		4	0.249	0.260	1.04 $\times$	0.153	0.150	0.98 $\times$
		16	0.259	0.332	1.28 $\times$	0.200	0.240	1.20 $\times$
2048	128	1	0.982	0.991	1.01 $\times$	0.381	0.442	1.16 $\times$
		4	0.997	1.020	1.02 $\times$	0.567	0.457	0.81 $\times$

Appendix

Table 5: Speculative verify against Batched-RPA : 32 query and 8 KV heads,

D=128

, 128-token pages, a window of

T

tokens per request. Time is ms per verify step, with JAX 0.11.0 and libtpu 0.0.44.

Batch	FlyWheel e4m3fn	FlyWheel e5m2	RPA-v3 e4m3fn
32	0.495	0.494	0.489
128	1.979	1.973	1.954

Appendix

Table 6: FP8 decode against RPA-v3 on one v6e chip: 32 query and 8 KV heads,

D=128

, 8K cached tokens per sequence. Time is ms per call.

		GDN			KDA
Step	Batch	Conv.	ReplaySSM	Speedup	Conv.	ReplaySSM	Speedup
Decode	16	0.101	0.057	1.77 $\times$	0.105	0.055	1.91 $\times$
	64	0.264	0.114	2.32 $\times$	0.267	0.123	2.17 $\times$
	128	0.478	0.207	2.31 $\times$	0.478	0.249	1.92 $\times$
	256	0.940	0.447	2.10 $\times$	0.944	0.500	1.89 $\times$
Verify-4	16	0.189	0.058	3.26 $\times$	0.205	0.065	3.15 $\times$

Appendix

Table 7: ReplaySSM decode on one v6e chip: 16 query/key heads, 32 value heads,

d_{k}=d_{v}=128

W=16

, BF16 buffered inputs. Conv. is FlyWheel’s conventional decode kernel. Time is ms per call; verify-4 checks 4 tokens per sequence.

Neural Network Hardware Acceleration

See all 63 papers

Systolic Linear Attention: Co-Designing Tile-Grouped Gated Delta Attention for TPUs
Nov 4, 2026Gagik Amirkhanyan, Nuojin ChengLinear Attention Attention Mechanisms
Google
Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU
Apr 16, 2026Jevin Jiang, Ying Chen, Blake A. Hechtman +2Efficient Transformer Inference AI Accelerator Inference
Google
ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers
Apr 26, 2026Chih-Chung Hsu, Xin-Di Ma, Wo-Ting Liao +1Efficient Transformer Inference Softmax Attention
Advanced Computer Vision Laboratory, National Yang Ming Chiao Tung University

Nov 4, 2026J/K move · Enter open · S save

Systolic Linear Attention: Co-Designing Tile-Grouped Gated Delta Attention for TPUs

Gagik Amirkhanyan, Nuojin Cheng

Google

Linear Attention Attention Mechanisms Neural Network Hardware Acceleration Gated Delta Networks

Open paper AskPDF arXiv

Gated-delta linear attention layers such as Gated DeltaNet (GDN), Kimi Delta Attention (KDA), and GDN-2 are tuned for GPUs, where small tensor-core tiles and fast vector registers make per-channel decay rescaling and sequential in-block solves cheap. On Google TPUs, which concentrate compute in wide

256 \times 256

systolic matrix units backed by a slower vector unit, porting these layers directly leaves half of each matrix tile empty and bottlenecks on vector operations. We present Systolic Linear Attention (SLA), which reshapes gated delta attention around the matrix unit by widening keys to

256

with grouped value heads, grouping decay gates into four

64

-channel tiles, and computing the chunk solve entirely with block matrix products in a single fused kernel. On TPU v7x, a full SLA-R4 layer runs

1.6\times

2.5\times

faster than the open-source Tokamax KDA kernel (

1.4\times

2.4\times

with exact chunk inversion), driven primarily by head sharing and kernel fusion. Across a

2{,}048

-chip TPU v7x pre-training study of sparse Mixture-of-Experts (MoE) models from

6.5\text{B}

63.1\text{B}

parameters, SLA trains

11\%

24\%

faster than KDA at similar perplexity, and reaches slightly lower perplexity than KDA in smaller matched-seed runs.

Abstract

Neural Network Hardware Acceleration

Explore similar work

Systolic Linear Attention: Co-Designing Tile-Grouped Gated Delta Attention for TPUs

Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU

ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers