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

Gagik Amirkhanyan; Nuojin Cheng

cs.LGNov 4, 2026

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

Authors: Gagik Amirkhanyan, Nuojin Cheng

Organizations: Google

Abstract

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}$ to $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.

Figures & tables

Figure 1: Overview of SLA. (a) Direct ports of GPU-tuned gated-delta layers (GDN [ 68 ] , KDA [ 71 ] , GDN-2 [ 25 ] ) run as an unfused pipeline on TPUs: intermediates round-trip through HBM, per-channel decay rescaling in KDA and GDN-2 and in-block triangular substitution run on the VPU, and

128

-wide states leave at least half of each

256\times 256

MXU tile padded. (b) SLA fuses stages 2–5 into one VMEM kernel. (c) A head with key and value widths

d_{k}=d_{v}=128

fills a quarter of an MXU tile, while SLA’s wide packed state fills whole tiles. (d) Decay gates are grouped into four

64

-channel tiles, so intra-chunk scores are four

64

-wide MXU matmuls with no VPU rescaling. (e) The chunk solve uses only block matrix products on the MXU.

Figure 2: Sequential state-scan cost on TPU v7x. A

64

-step scan of

(C,256)\times(256,256)

matmuls takes nearly constant time for

C\in\{32,\dots,256\}

, so a fixed per-step overhead, not matmul FLOPs, dominates the scan.

Accelerator	Cores	Array Tile	BF16 TF/s	8-bit TF/s ∗	VPU TF/s	MXU:VPU	HBM (GB)	HBM GB/s	Ridge
H100 SXM	132 SM	$16\times 16$	989	1,979	67.0	$14.8\times$	80	3,350	295
TPU v4	2	$128\times 128$	275	275	3.2	$85.9\times$	32	1,200	229
TPU v5p	2	$128\times 128$	459	918	5.4	$85.0\times$	95	2,765	166
TPU v6e	1	$256\times 256$	918	1,836	7.2	$127.5\times$	32	1,600	574
TPU v7x	2	$256\times 256$	2,307	4,614	18.0	$128.2\times$	192	7,373	313

Table 1: Documented hardware specifications per chip for NVIDIA H100 SXM and four Google TPU generations [ 32 , 18 , 6 ] . MXU:VPU is the ratio of peak BF16 matrix throughput to FP32 vector throughput; Ridge is the BF16 arithmetic-intensity ridge point (FLOPs/byte). ∗ 8-bit peak is FP8 on H100 and TPU v7x and INT8 on TPU v4–v6e.

Symbol	Meaning
$T,\;C,\;c$	Sequence length; chunk length ( $C=128$ ); chunk index $c\in\{1,\dots,T/C\}$ .
$t,j,s$	Token indices inside a chunk, $1\leq j\leq t\leq C$ (with chunk-entry boundary $0$ in $P_{t\leftarrow 0}$ ).
$D,\;H_{k},\;H_{v},\;R$	Model width; number of query-key heads; number of value heads; sharing ratio $R=H_{v}/H_{k}\in\{2,4\}$ .
$d_{k},\;d_{v},\;D_{v,\text{pack}},\;D_{v}$	Key width ( $d_{k}=256$ ); per-value-head width; packed value width per key head $D_{v,\text{pack}}=Rd_{v}$ ; total value width $D_{v}=H_{v}d_{v}$ .
$G,\;b,\;g$	Number of decay tiles ( $G=4$ ), tile width $b=d_{k}/G=64$ , tile index $g\in\{1,\dots,G\}$ .
$q_{t},k_{t}\in\mathbb{R}^{d_{k}}$ , $v_{t}\in\mathbb{R}^{D_{v,\text{pack}}}$	$L_{2}$ -normalized query and key of one key head; packed values of its $R$ value heads.

Table 2: Summary of notation used across the paper.

Design Axis	GDN [ 68 ]	KDA [ 71 ]	GDN-2 [ 25 ]	SLA ( SLA-R2 / SLA-R4 ) (Ours)	SPD-SLA-W (Section 4 )
Key/Value Heads	Unpacked $H_{k}=H_{v}$	Unpacked $H_{k}=H_{v}$	Unpacked $H_{k}=H_{v}$	Packed GVA $H_{k}=H_{v}/R$ ( $R\!\in\!\{2,4\}$ )	Packed GVA $H_{k}=H_{v}/R$
State Tile & Dtype	$S^{(h_{k})}\in\mathbb{R}^{128\times d_{v}}$ FP32 ( $50\%$ pad)	$S^{(h_{k})}\in\mathbb{R}^{128\times d_{v}}$ FP32 ( $50\%$ pad)	$S^{(h_{k})}\in\mathbb{R}^{d_{k}\times d_{v}}$ FP32	$S\in\mathbb{R}^{256\times D_{v,\text{pack}}}$ BF16 / Row-FP8 (no padding)	$S\in\mathbb{R}^{256\times D_{v,\text{pack}}}$ BF16 / Row-FP8
Forget Gate $F_{t}$	Scalar $\gamma_{t}I_{d_{k}}$ ( $G=1$ )	Channel $\operatorname{diag}(f_{t})$ ( $G=d_{k}$ )	Decoupled ch. ( $G=d_{k}$ )	Tile-grouped $\bigoplus_{g=1}^{4}\gamma_{t,g}I_{64}$ ( $G=4$ )	Tile-grouped $\bigoplus_{g=1}^{4}\gamma_{t,g}I_{64}$ ( $G=4$ )
Step Size $\beta_{t}$	$(0,1)$ per $H_{v}$	$(0,1)$ per $H_{v}$	Ch. vector per $H_{v}$	$(0,2)$ scalar per $H_{k}$ (neg. eig.)	Clamped $(0,(2-\varepsilon_{\beta})/k_{t}^{\top}\tilde{k}_{t}]$
Chunk Gram $M_{c}$	$1\times$ MXU + scalar mask	$16$ -row VPU pivots ( $\|\Lambda_{j,i}\|\!>\!88.7$ )	Decoupled ch. broadcast	$G\times[C,b]\!\times\![b,C]$ pure MXU ( $b\!\geq\!64$ , $0$ ovf.)	Oblique $k_{t}^{\top}P_{t\leftarrow j}\tilde{k}_{j}$ MXU
Chunk Inverse $A_{c}$	$H_{v}\times$ blocked, sequential in-block	$H_{v}\times$ blocked, sequential in-block ( $C\leq 64$ Tokamax)	$H_{v}\times$ blocked, sequential in-block	$H_{k}\times$ Exact Inv. (or $1$ -Step Neumann, $C=128$ )	$H_{k}\times$ MXU inverse + Polar-2+ ( $5$ MM)

Table 3: Side-by-side algebraic and systolic hardware comparison across prior GPU gated-delta recurrences and our SLA family.

Configuration	$\mu$	Loss @ $7{,}199$
KDA (per-channel decay, FP32 state)	–	$2.2925$
SLA, $R=2$ (two runs)	$0$	$2.2912$ / $2.2925$
SLA, $R=2$	$0.25$	$2.2955$
SLA, $R=2$ , input-dependent $\mu$	adaptive	$2.3216$
SLA, $R=4$ (two runs)	$0$	$2.3103$ / $2.3148$
SLA, $R=4$ $\dagger$	$0.25$	$\mathbf{2.2914}$

Table 4: Chunk-boundary momentum ablation:

20

-step rolling training cross-entropy at step

7{,}199

(single pass over Dolma 3.5,

12.5\text{B}

MoE,

128

TPU v4 chips, seed

42

; lower is better). † Trained on an

8{,}200

-step cosine schedule and evaluated at step

7{,}199

Figure 3: Detailed architectural diagram of the three mathematical pillars of SLA. (a) Pillar 1: Widening the key dimension to

d_{k}=256

and sharing each key head

h_{k}

and write gate

\beta_{t}\in(0,2)

across

R\in\{2,4\}

value heads forms a dense

256\times D_{v,\text{pack}}

state that fills whole

256\times 256

MXU tiles without padding while cutting QK projections, short convolutions, and chunk solves by

R\times

. (b) Pillar 2: Partitioning the

d_{k}=256

key channels into

G=4

contiguous tiles of

b=64

channels decomposes the intra-chunk coupling matrix

M_{c}

into four

64

-wide MXU matrix multiplications weighted by scalar decay masks

\Gamma_{g}

, eliminating VPU sub-block rescaling. (c) Pillar 3: The default Exact Inv. computes

(I_{C}+M_{c})^{-1}

by recursive block doubling on the MXU; the optional

1

-step Neumann inverse keeps only the first off-diagonal term of

I_{C}+M_{c}=D_{\text{blk}}(I_{C}+N)

with

16\times 16

diagonal blocks, where

N^{8}=0

C=128

. Both avoid sequential VPU substitution and use dense MXU matrix multiplications (Section 2.4.3 ).

Figure 4: Single-kernel VMEM execution and custom VJP backward pass in SLA. Top (Forward): Streamed

Q,K,V

and gates enter the

64\text{ MiB}

VMEM scratchpad once; a

3

-token ring buffer executes causal Conv1D and

L_{2}

normalization (Op 2), followed in-place by

4

-tile MXU Gram contractions (Op 3), MXU chunk inversion and WY factors (Ops 4–5), persistent

[256,D_{v,\text{pack}}]

state updates (Op 6), and chunk readout (Op 7a) before writing

O

to HBM. Bottom (Backward): During the reverse state scan, the custom VJP kernel recomputes intra-chunk matrices

M_{c},A_{c}

on the MXU directly inside VMEM, eliminating

[B,H_{k},T/C,C,C]

intermediate HBM writes for batch size

B

Stage Latency ( $D=2048,T=8\text{K}$ )	KDA (Unf.)	KDA (Fus.)	SLA-R4 (Unf.)	SLA-R4 (Fused)
Op 1: QKV + Gate Proj (MXU)	0.251	0.251	0.168	0.168
Op 2: Causal Conv4 + QK Norm	0.562	$\left.\begin{array}[]{@{}c@{}}\text{Fused}\\[0.77498pt] 3.104\end{array}\right\}$	0.341	$\left.\begin{array}[]{@{}c@{}}\text{Fused}\\[0.77498pt] 1.112\end{array}\right\}$
Op 3–5: Gram + WY Inversion	1.184		0.418
Op 6: Chunk State Scan ( $S_{c}$ )	1.248		0.614
Op 7a: Intra + State Readout	0.219		0.219
Op 7b–7c: Out Norm + Out Proj	0.255	0.255	0.245	0.245

Table 5: Seven-operator forward audit ( ms ) on TPU v7x (

D=2048

T=8{,}192

). Unfused columns report isolated per-stage forward medians; fused brackets report single-kernel Ops 2–7a latency.

Regime	Kernel or Full Layer ( $H_{k}\times d_{k}\times Rd_{v}$ , Chunk $C$ , State)	$8\text{K}$ Fwd	$8\text{K}$ Bwd	$8\text{K}$ F+B	$8\text{K}$ Spd.	$32\text{K}$ F+B	$32\text{K}$ Spd.	$32\text{K}$ Temp HBM
(a) Core Kernel $D=1792$ ( $63.1\text{B}$ MoE)	Open-Source MaxText Pallas KDA ( $8\times 256\times 512,C=128$ , FP32)	3.320	2.000	5.320	$0.63\times$	20.430	$0.61\times$	–
	Open-Source Tokamax Mosaic KDA ( $8\times 256\times 512,C=64$ , Base)	1.951	1.403	3.354	$1.00\times$	12.449	$1.00\times$	–
	SLA-R2 ( $4\times 256\times 1024,C=128$ , BF16)	1.960	0.905	2.865	$1.17\times$	10.153	$1.23\times$	–
	SLA-R4 ( $2\times 256\times 2048,C=128$ , BF16)	1.371	0.671	2.042	$1.64\times$	6.175	$2.02\times$	–
(b) Iso-Controls Core Kernel $D=2048$	Open-Source MaxText Pallas KDA ( $8\times 128\times 256,C=128$ , FP32)	2.668	1.580	4.248	$0.44\times$	16.455	$0.35\times$	–
	Open-Source Tokamax Mosaic KDA ( $8\times 128\times 256,C=64$ , Base)	1.072	0.781	1.853	$1.00\times$	5.829	$1.00\times$	–

Table 6: Measured TPU v7x (batch size

B=1

) kernel and full-layer latency ( ms ) ,

32\text{K}

temporary HBM (

8\text{K}

{\sim}0.26\times

), and F+B speedups normalized to Open-Source Tokamax Mosaic KDA Base (

1.00\times

) in each panel across (a) core kernels at

D=1792

, (b) iso-controlled core kernels at

D=2048

, and (c1–c2) full linear attention layers (Ops 1–7c, unified harness; Tokamax supports

C\leq 64

and

d_{k}\leq 256

); bold marks the SLA-R4 rows.

Figure 5: Forward speedup of SLA-R4 across TPU generations over a portable FP32 KDA baseline, with leave-one-out roofline predictions (

3.1\%

error). Not comparable to the Tokamax-normalized speedups in Table 6 .

TPU Generation	Array Tile	MXU:VPU (BF16)	KDA Measured (ms)	SLA-R4 Measured (ms)	Measured Speedup	Out-of-Fold LOO Pred. (Error)
TPU v4	$128\times 128$	$85.9\times$	$14.84$	$3.32$	$4.47\times$	$4.40\times$ ( $-1.6\%$ )
TPU v5p	$128\times 128$	$85.0\times$	$9.19$	$1.98$	$4.64\times$	$4.71\times$ ( $+1.5\%$ )
TPU v6e	$256\times 256$	$127.5\times$	$11.02$	$1.88$	$5.86\times$	$6.15\times$ ( $+5.0\%$ )
TPU v7x	$256\times 256$	$128.2\times$	$6.99$	$1.05$	$6.65\times$	$6.38\times$ ( $-4.1\%$ )

Table 7: Measured forward core-kernel latency ( ms ,

B=1,T=8{,}192,D=2048

) vs. portable JAX/Pallas KDA (

C=64

) and out-of-fold LOO Roofline predictions across four TPU generations.

Figure 6: Context parallelism and serving on TPU v7x. (a) Forward-backward latency of context-parallel scans across

8

–

128

TensorCores (Section 3.3 ). (b) Single-chip

32

-layer serving throughput and

131\text{K}

prefill latency (Table 8 ).

Regime	Architecture / Method ( $32$ L, $D=2048$ )	$8\text{K}$ TTFT	$32\text{K}$ TTFT	$131\text{K}$ TTFT	$32\text{K}$ Cache/Lyr
(a) $32$ L Prefill Full TTFT ( $B=1$ )	0L:1A Dense GQA ( $H_{q}=16$ , $H_{kv}=4$ KV heads)	$214.7\text{ ms}$	$761.1\text{ ms}$	$7{,}139\text{ ms}$ ( $1.00\times$ )	$64.0\text{ MiB}$ KV
	15L:1A MaxText KDA ( $8\times 128\times 256$ , FP32)	$265.7\text{ ms}$	$689.1\text{ ms}$	$2{,}664\text{ ms}$ ( $2.68\times$ )	$1.0\text{ MiB}$ FP32
	7L:1A SLA-R4 ( $2\times 256\times 1024$ , BF16)	$204.3\text{ ms}$	$439.8\text{ ms}$	$2{,}034\text{ ms}$ ( $3.51\times$ )	$1.0\text{ MiB}$ BF16
	15L:1A SLA-R4 ( $2\times 256\times 1024$ , FP8)	$203.6\text{ ms}$	$416.9\text{ ms}$	$1{,}669\text{ ms}$ ( $4.28\times$ )	$0.5\text{ MiB}$ FP8
Regime	Attention-Layer Serving ( tok/s/chip ) †	$8\text{K},B\!=\!128$	$32\text{K},B\!=\!64$	$131\text{K},B\!=\!32$	KV+State ( $131\text{K}$ )
(b) Serving Tree Verify ( $T_{\text{ver}}\in\{8,16\}$ )	0L:1A Dense GQA ( $T_{\text{ver}}=8$ )	$25.0\text{K}$	OOM	OOM	$256.0\text{ GiB}$

Table 8: Single-chip

32

-layer prefill latency (a) ,

32

-layer attention-layer serving throughput with speculative verification of

T_{\text{ver}}\in\{8,16\}

drafted tokens (b; OOM

>190\text{ GiB}

/chip), and context-parallel F+B scan latency across

P\in\{32,64,128\}

TensorCores (c) on TPU v7x (

D=2048

(a) $12.5\text{B}$ MoE CP=EP=128 F+B	$T=128\text{K}$ ( tok/s/ch )	$T=256\text{K}$ ( tok/s/ch )	$T=512\text{K}$ ( tok/s/ch )	$T=1.05\text{M}$ ( tok/s/ch )	$1\text{M}$ HBM / Mixer ICI
0L:1A Full-Attn Ring-CP	$1{,}740.4\text{ ms}$ ( $1{,}177$ )	$4{,}797.5\text{ ms}$ ( $854$ )	$10{,}712.1\text{ ms}$ ( $765$ )	$36{,}725.0\text{ ms}$ ( $446$ )	$25.8\text{ GiB}$ / $32{,}512\text{ MiB}$
7L:1A MaxText KDA Tree ( $G=128$ )	$1{,}232.2\text{ ms}$ ( $1{,}662$ )	$2{,}573.3\text{ ms}$ ( $1{,}592$ )	$2{,}704.5\text{ ms}$ ( $3{,}029$ )	$7{,}377.0\text{ ms}$ ( $2{,}221$ )	$25.73\text{ GiB}$ / $4{,}232\text{ MiB}$
7L:1A Iso-Param SLA Tree-CP R4	$\mathbf{1{,}212.9\text{ ms}}$ ( $\mathbf{1{,}689}$ )	$\mathbf{2{,}500.8\text{ ms}}$ ( $\mathbf{1{,}638}$ )	$\mathbf{2{,}542.6\text{ ms}}$ ( $\mathbf{3{,}222}$ )	$\mathbf{7{,}023.9\text{ ms}}$ ( $\mathbf{2{,}333}$ )	$25.92\text{ GiB}$ / $4{,}736\text{ MiB}$
15L:1A MaxText KDA Tree ( $G=128$ )	$1{,}221.4\text{ ms}$ ( $1{,}677$ )	$2{,}498.2\text{ ms}$ ( $1{,}639$ )	$2{,}249.8\text{ ms}$ ( $3{,}641$ )	$5{,}337.6\text{ ms}$ ( $3{,}070$ )	$25.49\text{ GiB}$ / $2{,}760\text{ MiB}$
15L:1A Ring-State LASP R2	$1{,}575.5\text{ ms}$ ( $1{,}299$ )	$2{,}744.8\text{ ms}$ ( $1{,}492$ )	$2{,}346.3\text{ ms}$ ( $3{,}491$ )	$5{,}204.6\text{ ms}$ ( $3{,}148$ )	$25.46\text{ GiB}$ / $17{,}392\text{ MiB}$
15L:1A SLA Tree-CP R2 NoMom / +Mom	$\mathbf{1{,}174/1{,}189\text{ ms}}$	$\mathbf{2{,}331/2{,}368\text{ ms}}$	$\mathbf{1{,}965/2{,}003\text{ ms}}$ ( $\mathbf{5.45\times}$ / $5.35\times$ )	$\mathbf{4{,}738/4{,}904\text{ ms}}$ ( $\mathbf{3{,}458/3{,}341}$ )	$\mathbf{25.45\text{ GiB}/2{,}992\text{ MiB}}$ ( $\mathbf{-83\%}$ )

Table 9: Measured

64

-chip TPU v7x ( 4x4x4 ,

128

TensorCores) end-to-end

16

-layer

12.5\text{B}

sparse MoE (

0.79\text{B}

active,

256

experts

k=8

) (a) fixed- CP=EP=128 F+B latency and tok/s/chip across

T=128\text{K}

–

1.05\text{M}

(

T_{\text{loc}}=T/128

; at

T=512\text{K}

T_{\text{loc}}=4{,}096

selects a faster XLA EP=128 MoE ragged-all-to-all tile schedule across all linear hybrids relative to

T=256\text{K}

), (b) EP=128 cube verification throughput (

B_{\text{chip}}=64

, comparing Tree-N1 [= Tree-Neumann-1 ,

T_{\text{ver}}=16

] and FP8Fused seq_scan_t16 ), and (c) real-checkpoint lossless speculative acceptance

\bar{\tau}

(

K=16

N=64

, greedy [stochastic] ).

Architecture / Scaling Parameter	Rung 1	Rung 2	Rung 3	Rung 4
Cohort Label (Active / Total)	$0.35\text{B}/6.5\text{B}^{\ast}$	$0.79\text{B}/12.5\text{B}$	$1.80\text{B}/31.5\text{B}^{\ast}$	$3.74\text{B}/63.1\text{B}$
Exact Total Params (KDA / SLA-R2 )	$7.58\text{B}/7.57\text{B}$	$12.49\text{B}/12.49\text{B}$	$30.58\text{B}/30.52\text{B}$	$63.13\text{B}/63.02\text{B}$
Exact Active: incl. $2VD$ ( $1VD$ FLOP) ∗	$0.44\text{B}$ ( $0.34\text{B}$ )	$0.79\text{B}$ ( $0.69\text{B}$ )	$1.91\text{B}$ ( $1.73\text{B}$ )	$3.74\text{B}$ ( $3.56\text{B}$ )
Decoder Depth $L$ & Hybrid Cycle	$10$ ( 4L : $8{+}2$ )	$16$ ( 7L : $14{+}2$ ; 15L )	$15$ ( 4L : $12{+}3$ )	$30$ ( 4L : $24{+}6$ ; 15L )
Model Width $D$ ( $4\times\operatorname{RMSNorm}$ /block)	$1024$ ( $4{\times}256$ )	$1024$ ( $4{\times}256$ )	$1792$ ( $7{\times}256$ )	$1792$ ( $7{\times}256$ )
Embeddings + Untied LM Head ( $2VD$ )	$205.5\text{M}$ ( $\sqrt{D}{+}\text{Norm}$ )	$205.5\text{M}$ ( $\sqrt{D}{+}\text{Norm}$ )	$359.7\text{M}$ ( $\sqrt{D}{+}\text{Norm}$ )	$359.7\text{M}$ ( $\sqrt{D}{+}\text{Norm}$ )

Table 10: Exact model architecture and multi-scale parameterization across our four hybrid LatentMoE pre-training rungs (

V=100{,}352

untied vocabulary,

E=256

routed experts on layers

1,\dots,L-1

, dense SwiGLU on layer

0

Figure 7: (a) Per-step XProf breakdown for

63.1\text{B}

MoE on

64

TPU v7x chips across KDA Base , SPD-SLA-W (R2) ( SLA-R2 with optional in-chunk key whitening, Section 4 ), and SLA-R4 (Table 13 ); most of SLA’s saving is in linear attention ops. (b) Held-out Dolma tail perplexity from

4\text{K}

32\text{K}

for

12.5\text{B}

MoEs (Table 11 a; seed s0 ). (c) Zero-shot

8\text{K}\to 64\text{K}

RULER HotpotQA for KDA (raw and clamped) and SLA (

60.4\text{B}

-token curves: seed s0 ;

15.1\text{B}

-token curves: mean of s0 and s1 ).

Cohort	Architecture ( $7{,}200$ Steps)	Step (ms)	$11$ -Set PPL	$8\text{K}$ Tail	$32\text{K}$ Tail	$32\text{K}$ $8$ -N	$64\text{K}$ SQuAD	$64\text{K}$ Hotpot
(a) v7x $60.4\text{B}$ Tok ( $3.8\times$ Ch., $8.4\text{M}$ /st.)	0L:1A Full-Attn GQA+RoPE ( s0 )	2,313.8	16.128	10.57	10.48	$12.1\%$	$58.6\%$	$53.1\%$
	7L:1A MaxText KDA Base ( $n=3$ ) ∗	$2{,}225\pm 29$	$15.602\pm 0.582$	10.46	$12.45\pm 1.11$	$24.6\%$	$54.9\%$	$50.3\%$
	$+$ Zero-Retrain $\tau_{\max}=8\text{K}$ Clamp ( s0 )	–	16.168	10.53	10.26	$57.8\%$	$\mathbf{78.1\%}$	$75.0\%$
	7L:1A KDA + TrainMaxTau=8K ( $n=2$ )	$2{,}251\pm 11$	$15.800\pm 0.332$	10.46	$10.75\pm 0.21$	$49.6\%$	$75.8\%$	$75.4\%$
	7L:1A SLA-R2 (BF16, $n=3$ ) ∗	$1{,}919\pm 37$	$15.580\pm 0.508$	10.46	$\mathbf{10.19\pm 0.10}$	$56.6\%$	$75.5\%$	$\mathbf{77.9\%}$
	7L:1A SLA-R4 (Iso-Param, $n=3$ ) ∗†	$\mathbf{1{,}800\pm 39}^{\dagger}$	$15.664\pm 0.414$	10.57	$10.27\pm 0.04$	$56.1\%$	$74.7\%$	$74.5\%$

Table 11: Measured

12.5\text{B}

MoE (

0.79\text{B}

active,

16

D=1024,D_{v}=2048,T_{\text{train}}=8{,}192

) pre-training across (a)

60.40\text{B}

-token (

3.8\times

Chinchilla [ 27 ] ,

8.39\text{M}

tok/step) v7x arms and (b)

15.10\text{B}

-token (

1.0\times

Chinchilla,

2.10\text{M}

/step) and

{\sim}90

-min Iso-Wall ablations (

n\in\{2,3\}

random seeds where noted;

4

-choice

64\text{K}

RULER QA,

N=128

\text{SE}=4\text{ pts}

; full controls in Tables 15 – 16 ; bold: best within each cohort at equal tokens, with perplexity compared on matched seeds s0,s1 in (b); GDN-2 uses our unfused TPU port; seeds s0 / s1 / s2 use different data offsets, so perplexities are comparable only on matched seeds).

$63.1\text{B}$ MoE ( $120.8\text{B}$ Tok, $512$ Chips, NoPE $8\text{K}$ )	Step (ms)	tok/s/ch	$11$ -Set PPL	$32\text{K}$ Tail	$32\text{K}$ $8$ -N	$64\text{K}$ SQuAD	$64\text{K}$ Hotpot
4L:1A MaxText KDA Base ( $8\times 256\times 512$ , FP32)	2,165.1	15,135	12.596	353.4	$15.6\%$	$27.7\%$	$27.0\%$
$+$ Zero-Retrain $\tau_{\max}=8\text{K}$ Clamp ( $0$ retr.)	–	–	12.596	7.70	$48.4\%$	$76.6\%^{\ast}$	$82.0\%^{\ast}$
$+\tau_{\max}=8\text{K}$ Clamp $+$ SSMax ( $\alpha=0.25$ )	–	–	12.596	8.04	$74.2\%$	$\mathbf{90.6\%}$	$\mathbf{86.5\%}$
4L:1A SLA-R2 ( $4\times 256\times 1024$ , BF16, $0.5\times$ Bytes)	1,952.9	16,779	12.584	13.78	$18.6\%$	$55.7\%$	$56.8\%$
$+$ Zero-Clamp SLA-R2 $+$ SSMax ( $\alpha=0.5$ )	–	–	12.584	8.45	$\mathbf{84.4\%}$	–	–
4L:1A SLA-R4 ( $2\times 256\times 2048$ , BF16, $0.5\times$ Bytes) †	1,878.3	17,445	12.454	20.47	$18.0\%$	$51.6\%$	$45.1\%$

Table 12: Measured

512

-chip TPU v7x

63.1\text{B}

MoE (

3.74\text{B}

active,

8\times 8\times 8

16.78\text{M}

tok/step) at

120.80\text{B}

tokens (

\approx 1.6\times

Chinchilla, NoPE

T_{\text{train}}=8{,}192

; SSMax [ 40 ] query scale

\alpha\in\{0.25,0.5\}

, Section 3.5 ;

64\text{K}

RULER QA cand4_acc at

N=512

{}^{\ast}N=128

; full

64

-chip XProf and RoPE ablations in Tables 15 , 13 , and 17 ).

HLO Category ( xplane Measured)	4L:1A KDA Base	4L:1A SPD-SLA-W (R2)	4L:1A SLA-R4	$\Delta$ (SLA-R4 vs. KDA)
Pallas Linear Attention custom-call	$703.0\text{ ms}$	$597.2\text{ ms}$	$552.0\text{ ms}$	$-151.0\text{ ms}$ ( $-21.5\%$ )
VPU Gate / Norm loop fusion	$190.4\text{ ms}$	$156.8\text{ ms}$	$143.7\text{ ms}$	$-46.7\text{ ms}$ ( $-24.5\%$ )
Causal Short convolution fusion	$220.0\text{ ms}$	$207.4\text{ ms}$	$199.0\text{ ms}$	$-21.0\text{ ms}$ ( $-9.5\%$ )
ICI Collective async-done Wait	$366.6\text{ ms}$	$354.1\text{ ms}$	$343.1\text{ ms}$	$-23.5\text{ ms}$ ( $-6.4\%$ )
MoE Expert ragged-dot + Dense Matmul	$245.8\text{ ms}$	$266.8\text{ ms}$	$237.2\text{ ms}$	$-8.6\text{ ms}$ ( $-3.5\%$ )
Total Steady-State Step Time	$\mathbf{1{,}725.8\text{ ms}}$	$\mathbf{1{,}582.3\text{ ms}}$	$\mathbf{1{,}475.0\text{ ms}}$	$\mathbf{-250.8\text{ ms}}$ ( $\mathbf{1.170\times}$ )

Table 13: Hardware XProf per-step HLO category breakdown ( ms/step ) on

64

-chip TPU v7x

63.1\text{B}

MoE ( FSDP=2, EP=64 ,

2.10\text{M}

tok/step).

$12.5\text{B}$ GA=4 Arm ( $3$ profiled steps)	Step Mean (ms)	linear_mixer _pallas	short_conv _fusion	Mixer + Conv	full_attn _splash	all_gather _rs	moe_ragged _a2a
7L:1A MaxText KDA (default AR)	2,118.9	508.1	91.7	599.8	69.8	191.9	264.7
7L:1A MaxText KDA + defer_ar	2,119.7 ( $1.00\times$ )	528.4	91.6	620.0	70.0	191.6	234.3
7L:1A SLA (Fused) R2+Mom + defer_ar	1,797.4 ( $1.18\times$ )	340.9 ( $-187.5$ )	65.4	406.3 ( $-213.7$ )	69.9	152.5	282.5
7L:1A Iso-Param SLA (Fused) R4 + defer_ar	1,686.0 ( $1.26\times$ )	269.1 ( $1.96\times$ )	61.7	330.8 ( $-289.2$ )	69.7	152.6	265.3

Table 14: Symmetric

4

-arm GA=4 (

8.39\text{M}

tok/step) hardware XProf leaf-op breakdown ( ms/step ) on

64

-chip TPU v7x

12.5\text{B}

MoE ( EP=128 ).

Scale, Horizon, & Trained Checkpoint	$G$	Within-Tile Std	Unclamped $32\text{K}$ Tail / $8$ -N	Intervened $32\text{K}$ Tail / $8$ -N (Proj. $G{=}4$ )	Unclamped $64\text{K}$ SQuAD / Hotpot	Intervened $64\text{K}$ QA (Proj. $G{=}4$ / Note)
$6.5\text{B}$ 4L:1A MaxText KDA ( $15.1\text{B}$ )	128	4.41	$14.64/14.1\%$	$\mathbf{11.77/69.9\%}$ ( $70.47$ )	$53.1\%/42.2\%$	$\mathbf{75.0\%/79.7\%}$ ( $24.2/39.8\%$ )
$6.5\text{B}$ 4L:1A KDA + TrainMaxTau ( $15.1\text{B}$ )	128	$\leq\tau_{\max}$	$\mathbf{12.36/48.4\%}$	$\mathbf{12.36/48.4\%}$ (co-capped)	$\mathbf{69.5\%/73.4\%}$	$\mathbf{69.5\%/73.4\%}$ ( $68.8/70.3\%$ $N{=}512$ )
$6.5\text{B}$ 4L:1A Iso-Param SLA-R4 ( $15.1\text{B}$ )	4	0.00	$\mathbf{11.94/48.8\%}$	$\mathbf{11.94/48.8\%}$ ( $11.94$ )	$\mathbf{67.2\%/76.6\%}$	$\mathbf{67.2\%/76.6\%}$ (invariant)
$12.5\text{B}$ 7L:1A MaxText KDA ( $60.4\text{B}$ , s0 )	128	4.58	$12.35/15.2\%$	$\mathbf{10.26/57.8\%}$ ( $57.85$ )	$47.7\%/44.5\%$	$\mathbf{78.1\%/75.0\%}$ ( $28.1/45.3\%$ )
$12.5\text{B}$ 7L:1A KDA + TrainMaxTau ( $60.4\text{B}$ , $n=2$ )	128	$\leq\tau_{\max}$	$\mathbf{10.75/49.6\%}$	$\mathbf{10.75/49.6\%}$ (co-capped)	$\mathbf{74.5\%/72.8\%}$ ( $N{=}512$ )	$\mathbf{74.5\%/72.8\%}$ ( $N{=}512$ )
$12.5\text{B}$ 15L:1A MaxText KDA ( $60.4\text{B}$ )	128	4.97	$10.48/15.2\%$	$\mathbf{9.66/42.2\%}$ ( $63.30$ )	$59.4\%/61.7\%$	$\mathbf{73.4\%/78.9\%}$ ( $25.8/34.4\%$ )

Table 17: Multi-scale causal

\tau_{\text{eff}}

dispersion, zero-retraining

\tau_{\max}=8{,}192

inference clamping, live TrainMaxTau=8K saturation, post-hoc

G=4

projection, and NoPE SSMax (

\alpha=0.5

) on

32\text{K}

Tail PPL,

32\text{K}

8

-N, and

64\text{K}

RULER QA (

N=128

unless noted

N{=}512

). The clamp evaluator ( a3_causal_clamp_eval ) scores its own held-out Dolma slice and paired

N=128

QA subset, so absolute Tail PPL and QA values differ from the standalone scorers of Tables 11 , 12 , 15 , and 16 (by

1

–

2

QA items,

0.8

–

1.6

pts) while within-row intervention deltas remain exact.

$[4,256,512]$ Exact Inv. FP32 Arm ( $\mu=0$ )	$G$ ( $b$ )	Within-Tile Std	s0,s1 Mean PPL ( $n{=}3$ Mean) ∗	Paired $\Delta\text{PPL}$ vs. $G=1$ ( s0..s2 )	$8\text{K}\to 32\text{K}$ Tail	$32\text{K}$ $4$ -N / $8$ -N	norm9	$N=512$ $64\text{K}$ SQuAD / Hotpot
Full-Tile Scalar Decay ( $n=3$ )	1 ( $256$ )	0.00	$17.946$ ( $17.689$ )	$0.000$ (ref)	$11.58\to\mathbf{11.26}$	$\mathbf{70.4\%/59.0\%}$	$48.30\%$	$\mathbf{67.8\pm 1.6\%/71.2\pm 1.5\%}$
Systolic Sub-Tile Default ( $n=3$ )	4 ( $64$ )	0.00	$17.749$ ( $17.538$ )	$\mathbf{-0.151}$ ( $-0.34,-0.06,-0.06$ )	$11.65\to\mathbf{11.41}$	$58.6\%/\mathbf{49.2\%}$	$48.10\%$	$\mathbf{64.5\pm 4.3\%/68.9\pm 3.3\%}^{\ast}$
Intermediate Sub-Tile ( $n=3$ )	16 ( $16$ )	4.34	$17.873$ ( $17.626$ )	$-0.064$ ( $-0.14,-0.00,-0.05$ )	$11.75\to 12.12\pm 0.72$	$41.4\%/27.1\%$	$48.35\%$	$60.3\pm 4.2\%/61.9\pm 5.4\%$
Fine Sub-Tile Decay ( $n=3$ )	64 ( $4$ )	5.05	$17.748$ ( $17.523$ )	$-0.167$ ( $-0.36,-0.04,-0.11$ )	$11.53\to 12.94\pm 0.23$	$27.7\%/19.0\%$	$48.61\%$	$53.5\pm 2.2\%/50.8\pm 1.8\%$
Per-Channel Wide-Key ( $n=2$ PPL, $3$ OOD)	256 ( $1$ )	4.78	$\mathbf{17.694}$ ( s0,s1 )	$\mathbf{-0.251}$ ( $-0.36,-0.14$ )	$11.53\to 12.86\pm 1.13$	$41.0\%/25.1\%$	$48.54\%$	$56.8\%/58.9\%^{\dagger}$

Table 18: Unconfounded

5

-point

[4,256,512]

decay-granularity sweep on the identical chunk_kda Exact Inv. FP32 kernel (

\mu=0

12.5\text{B}

MoE, 7L:1A ,

15.10\text{B}

tok), reporting matched s0,s1

11

-set PPL (valid across all

5

granularities), paired per-seed

\Delta\text{PPL}

vs.

G=1

8\text{K}\to 32\text{K}

Tail PPL,

32\text{K}

4/8

-N, lm-eval norm9 , and

N=512

(

1\text{ SE}=2.0

pts)

64\text{K}

RULER QA.

Figure 8: Polar-2+ minimax contraction, overwrite recall, and checkpoint retention audit. (a) Scalar inverse-square-root contraction

\lambda z_{2}(\lambda)^{2}

on normalized eigenvalue

\lambda\in[0.02,1.0]

(Proposition 4.2 ): unscaled polynomial factor

\alpha=1.000

hits

1.0492>1

, whereas Polar-2+ (safety factor

\alpha_{\text{safe}}=0.968

) bounds the peak to

0.9831<1

, keeping minimum eigenvalue

\lambda_{\min}(\Sigma_{c})\geq 0.0388

of key preconditioner

\Sigma_{c}

across

4{,}096

BF16 chunks (

524\text{K}

tok). (b)

8

-seed Hadamard overwrite accuracy (overwrite multiplicity

r=3

) at memory load

\rho_{\text{eff}}=N/d_{k}=0.25

(

N=16

stored keys,

d_{k}=64

). (c) Trained

63.1\text{B}

( 4L:1A ) and

12.5\text{B}

( 15L:1A ) checkpoint load

\rho_{\text{eff}}

Whitening / Second-Order Solver	Ops	Fwd (ms)	F+B (ms)	$\kappa_{\text{post}}$	$\lambda_{\min}$ ( $4{,}096$ c BF16)
MesaNet Gated CG-30	VPU+MXU	$3.050$	$27.640$	$1.02$	– (CG)
Panel-16 Cholesky (VPU)	Tri. Solve	$1.412$	$2.385$	$1.00$	$+0.0148$ (SPD)
Coupled Newton-Schulz ( NS-5 )	$10$ MM	$0.708$	$1.024$	$1.04$	$+0.0294$ (SPD)
Polar-2 Unscaled ( $\alpha=1.000$ )	$5$ MM	$0.615$	$0.895$	$1.102$	$-0.0412$ (Indef.)
Polar-2+ Safe ( $\alpha=0.968$ )	$\mathbf{5\text{ MM}}$	$0.617$	$0.897$	$1.102$	$+0.0388$ (SPD)

Table 19: Measured TPU v7x chunk-whitening latency (

64

chunks

=8{,}192

tokens,

C=128,H_{k}=4,d_{k}=256

; MesaNet CG-10 :

1.74/11.16\text{ ms}

), exact conditioned Gram

\kappa_{\text{post}}

(

\leq 1.123

in BF16), and

4{,}096

-chunk (

524\text{K}

-token,

\rho_{p}=0

) BF16 Woodbury stress-test

\lambda_{\min}

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Mechanism / Update Rule	Axis-Aligned (FP32)	Hadamard-Rot. (FP32)	Write-Once BF16 [FP32]	$r=2$ Overwrite BF16 [FP32]	$r=3$ Overwrite BF16 [FP32]	Mod- $2$ Parity ( $\kappa_{\text{forget}}$ )
Mamba-2 / SSD (Scalar Hebbian)	$2.14\text{ dB}$	$2.11\text{ dB}$	$1.85$ [ $2.11$ ]	$1.72$ [ $1.89$ ]	$1.40$ [ $1.52$ ]	$50.2\%$ ( $1.00$ )
KDA Real Channel Decay ( $\beta\in(0,1)$ )	$5.31\text{ dB}$	$5.29\text{ dB}$	$-0.84$ [ $-0.91$ ]	$0.07$ [ $0.26$ ]	$60.97$ [ $>140$ ]	$51.2\%$ ( $1.1\times 10^{7}$ )
Diagonal Preconditioned Delta [ 62 ]	$22.41\text{ dB}$	$6.12\text{ dB}$	$4.92$ [ $6.12$ ]	$16.80$ [ $19.05$ ]	$18.10$ [ $21.40$ ]	$51.4\%$ ( $1.1\times 10^{7}$ )
MesaNet Unforgetted ( $\gamma=1.00$ ) [ 64 ]	$95.08\text{ dB}$	$95.06\text{ dB}$	$\mathbf{10.63}$ [ $22.79$ ]	$2.67$ [ $3.23$ ]	$1.90$ [ $1.73$ ]	$50.8\%$ ( $1.00$ )
MesaNet Gated ( $\gamma=0.98$ )	$91.40\text{ dB}$	$91.38\text{ dB}$	$9.54$ [ $22.56$ ]	$7.32$ [ $8.35$ ]	$4.13$ [ $4.07$ ]	$51.0\%$ ( $1.00$ )
MesaNet Aggressively Gated ( $\gamma=0.80$ )	$74.20\text{ dB}$	$74.18\text{ dB}$	$5.78$ [ $26.09$ ]	$47.68$ [ $89.36$ ]	$29.53$ [ $30.22$ ]	$51.1\%$ ( $1.00$ )

Appendix

Table 20: Measured associative recall SNR ( dB , reported as BF16 [FP32] ) and Mod-

2

parity accuracy across linear attention update rules (

d_{k}=256,C=128

8

seeds,

\kappa_{2}(\mathcal{G})\approx 48

);

\kappa_{\text{forget}}=\kappa_{2}(F_{c})

is the condition number of the chunk forget gate.

Linear Attention

Linear attention reduces sequence-length costs by replacing pairwise token interactions with kernel features, recurrent states, or compact memory.

See all 140 papers

FlyWheel: Attention Kernels Built for the TPU
Nov 1, 2026Zishuo Bao, Zeshen Zhang, Linkun Chen +1Neural Network Hardware Acceleration LLM Inference Acceleration
New York University · New York University Shanghai · Google
Exact Linear Attention
May 13, 2026Weinuo OuEfficient Transformer Inference Self-Attention
Wuyi University
When Good Enough Is Optimal: Multiplication-Only Matrix Inversion Approximation for Quantized Gated DeltaNet
Jun 4, 2026Luoming Zhang, Yuwei Ren, Kui Zhang +7Efficient Transformer Inference LLM Inference Acceleration
Qualcomm AI Research, an initiative of Qualcomm Technologies, Inc.

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

FlyWheel: Attention Kernels Built for the TPU

Zishuo Bao, Zeshen Zhang, Linkun Chen, Yucheng Lu

New York University · New York University Shanghai · Google

Neural Network Hardware Acceleration LLM Inference Acceleration Softmax Attention

Open paper AskPDF arXiv

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.

Abstract

Linear Attention

Explore similar work

FlyWheel: Attention Kernels Built for the TPU

Exact Linear Attention

When Good Enough Is Optimal: Multiplication-Only Matrix Inversion Approximation for Quantized Gated DeltaNet