FLAT brings linear-attention training kernels to TPUs

Researchers have released FLAT, an open-source set of Pallas kernels for training Linear Attention, Mamba-2 and Gated Linear Attention on Google TPUs. The work includes forward and backward passes plus the short causal convolution used before the recurrence.

The useful engineering story is how the kernels fit the hardware. In their technical write-up, the authors describe packing several attention heads into each matrix multiplication, moving prefix sums onto the matrix unit, and keeping recurrent state in on-chip memory throughout the forward pass.

Fast kernels alone were insufficient. Their first implementation still spent 1.8 ms per full-layer step copying between tensor layouts, plus another 1.1 ms splitting a fused projection and concatenating its gradient. The shipped version reads query, key and value directly from that projection and writes a fused gradient.

The public API exposes this as linear_attention_packed_jax. The repository requires Python 3.12 or later, uses Apache 2.0 licensing and provides a plain JAX reference backend alongside the TPU implementation.

The authors report 3.08× faster combined forward-and-backward execution for GLA on one TPU v6e chip than their FLA baseline on an H100. That is a kernel comparison. Full-layer training steps improved by 2.12×–2.75× against their JAX baseline, using batch size 8, sequence length 4,096 and bf16 inputs.

The remaining scope matters: DeltaNet, Gated DeltaNet and Kimi Delta Attention kernels are future work, as are tests on other TPU generations and sequence sharding across chips. The release supplies training building blocks; its measurements do not establish a speedup for an entire model.