SLA：为 TPU 重新设计的门控 Delta 线性注意力
Systolic Linear Attention: Co-Designing Tile-Grouped Gated Delta Attention for TPUs
KDA 这些线性注意力层都是为 GPU 调优的，搬到 TPU 会浪费一半算力，这篇把它们按 TPU 矩阵单元重新设计，训练快最多 24%，做 TPU 训练的可以看看。
论文提出 Systolic Linear Attention（SLA），针对 Google TPU 的 256×256 脉动矩阵单元重新设计 Gated DeltaNet、Kimi Delta Attention（KDA）等门控 Delta 线性注意力层，包括将 key 维度扩到 256、把 decay gate 分成四个 64 通道 tile 并用单个融合 kernel 完成块矩阵求解。在 TPU v7x 上，SLA-R4 层比开源 Tokamax KDA kernel 快 1.6–2.5 倍（精确 chunk 求逆时为 1.4–2.4 倍）。在 2048 芯 TPU v7x 预训练实验中，6.5B 到 63.1B 参数的稀疏 MoE 模型用 SLA 训练比 KDA 快 11%–24%，困惑度相近，小规模同 seed 运行中困惑度略低于 KDA。
Systolic Linear Attention: Co-Designing Tile-Grouped Gated Delta Attention for TPUs
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 \textbf{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 \texttt{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.
