FLAT

FLAT: fast, open-source TPU training kernels for linear attention and Mamba-2. TPU kernels face different bottlenecks than GPU kernels.