What is Systolic Linear Attention?

SLA is a redesign of gated-delta linear attention for TPUs. It widens keys to 256 with grouped value heads, groups decay gates into four 64-channel tiles, and computes the chunk solve with block matrix products in one fused kernel.

Why does a GPU-tuned linear-attention layer run poorly on a TPU?

A TPU v7x matrix unit is a 256×256 systolic array, so a 128-wide head leaves at least half of each tile padded. Per-channel decay and the sequential triangular solve also land on the slower vector unit.

How much faster is SLA?

The paper reports a full SLA-R4 layer 1.6×–2.5× faster than the open-source Tokamax KDA kernel on TPU v7x, and 11%–24% faster training than KDA at similar perplexity in a 2,048-chip study. The comparisons also change the model, not only the kernel.

Does this apply to GPUs?

The paper measures TPUs only. The general lesson, that a layer's tile shapes should match the hardware matrix unit, applies to any accelerator, but the specific SLA numbers do not transfer to GPUs.

GPU· 2026-11-06

Co-design linear attention for TPU systolic arrays — Systolic Linear Attention — What does it mean?

Jargon

Linear attention: An attention variant that keeps a fixed-size state matrix instead of a KV cache that grows with every token, so the cost per token stays constant.
Gated delta rule (GDN, KDA, GDN-2): A linear-attention update that compares what the state already predicts for the current key with the new value, then writes back the error scaled by a write gate (the delta rule), and also lets old content fade through a decay gate. KDA is Kimi Delta Attention.
Decay gate: A number between 0 and 1 that scales down what the state remembers. KDA and GDN-2 use one decay value per channel, which is more expressive but costs more arithmetic.
Systolic array: A grid of multiply-accumulate cells where data flows from neighbor to neighbor, so a matrix product runs with little memory traffic. The cost is that the grid works at one fixed size.
MXU: The matrix multiply unit of a TPU, a 256×256 systolic array on TPU v7x. It is the GPU tensor core's counterpart, but its tile is much larger.
VPU: The TPU vector unit, which handles element-wise work such as rescaling by a decay value. The paper describes it as slower than the MXU.
VMEM: The fast on-chip scratchpad of a TPU (64 MiB in the paper). A fused kernel keeps its intermediate results here instead of writing them to HBM, the chip's main memory.
Chunk solve: Linear attention processes tokens in chunks. Inside each chunk it must solve a small triangular system, which a GPU kernel does row by row. SLA rewrites it as block matrix products.

The news. On November 4, 2026, Google researchers posted Systolic Linear Attention (SLA). On a 2,048-chip TPU v7x pre-training study of sparse mixture-of-experts models from 6.5B to 63.1B parameters, SLA trained 11%–24% faster than KDA at similar perplexity. A full SLA layer ran 1.6×–2.5× faster than the open-source Tokamax KDA kernel. Read the paper →

Picture the matrix unit as a baking sheet that is always the same size. The sheet bakes in one pass whether it holds a full tray or a single small batch, so a mostly empty sheet wastes oven time. A TPU v7x matrix unit is a 256×256 sheet. GPU tensor cores are far smaller tiles, which is why GPU-tuned linear-attention layers, with 128-wide keys and values per head, fit them well. Put the same 128-wide head on a TPU and it covers about a quarter of the sheet. This is the same rule as why matrix dimensions must align to the hardware tile, pushed to a larger tile.

Linear attention itself is simple to state: instead of a KV cache that grows, it keeps one state matrix and updates it for every token. The gated delta layers (Gated DeltaNet, KDA, GDN-2) add two refinements, a delta-rule overwrite (erase the old value stored for a key, then write the new one) and a decay gate, and on a GPU the per-channel decay and the in-chunk solve they need are cheap. On a TPU the paper identifies three bottlenecks: narrow states that leave tiles padded, per-channel decay that runs on the slow vector unit, and a sequential triangular solve that also runs there. SLA answers each one in turn, as the table shows.

Bottleneck on a TPU	What a direct GPU port does	SLA's change
Narrow state	128-wide keys and values fill only a fraction of each 256×256 tile	Widen keys to 256 and let four value heads share them, so the packed state fills whole tiles
Per-channel decay	Every channel is rescaled separately on the vector unit	Group the decay gates into four 64-channel tiles, so intra-chunk scores become four 64-wide matrix products
Sequential solve	The in-chunk triangular solve runs step by step on the vector unit	Express the solve as block matrix products (block-doubling or a block Neumann series) on the matrix unit

The last piece is fusion. Once every stage of the chunk computation is a matrix product, the whole forward pass can run as one kernel that keeps its intermediate results in the 64 MiB VMEM scratchpad, the chip's fast on-chip memory. The backward pass recomputes two intermediate matrices instead of storing them in HBM. Skipping those round trips is the same idea as operator fusion: fewer trips to slow memory, more time inside the compute unit. It also echoes the data reuse that makes tiling work.

The paper's stage breakdown shows how the pieces combine. Fusing the GPU-shaped KDA layer gains only 1.03×. Reshaping the layer without fusing it gives 1.79×. Doing both gives 2.38×. Two numbers explain why. First, the tile arithmetic: a 256×256 tile has 65,536 cells and a 128×128 state has 16,384, which is 25% (this is geometry, assuming the compiler pads the state to the full tile). Second, the ratio of the last two measurements: 2.38 divided by 1.79 is about 1.33 (a figure derived here from the two reported totals, not a separately reported benchmark), so the same fusion that added 3% on the old shape adds about 33% once the layer is reshaped. Fusion only pays when the work inside it is already matrix-shaped.

Read the results with their limits in mind. The paper notes that its comparisons change the model as well as the kernel: SLA uses fewer key heads, a BF16 state, and grouped decay. Its abstract credits the kernel speedup primarily to head sharing and fusion, and the authors say they did not build a fused KDA or GDN kernel with the same 256-wide, 4:1 head sharing, so the tile-grouped gate's own contribution is not separated out. Without 4:1 sharing, the 256-wide SLA variants ran at 0.67×–0.77× of the Tokamax baseline. All measurements are on TPUs; the paper makes no claim about GPUs. On an identical portable kernel stack, SLA's forward speedup over a portable FP32 KDA baseline grows from 4.5× on TPU v4 to 6.7× on TPU v7x, which tracks the widening gap between matrix-unit and vector-unit throughput (a different baseline from the Tokamax numbers above).

Goes deeper in: GPU & CUDA → Tensor Cores & Mixed Precision → Dimension Alignment

Frequently Asked Questions

Check what you knowMap your AI & GPU knowledge across every track — free, role-based