Systolic Linear Attention：面向 TPU 的分块分组门控 Delta 注意力协同设计

arXiv:2611.06828v1 公告类型：新论文 摘要：Gated DeltaNet (GDN)、Kimi Delta Attention (KDA) 和 GDN-2 等门控 Delta 线性注意力层针对 GPU 进行了优化；GPU 的小型张量核心分块和高速向量寄存器使逐通道衰减重缩放与块内顺序求解的开销较低。Google TPU 将计算集中在宽大的 $256 \times 256$ 脉动矩阵单元上，并由较慢的向量单元提供支持，因此直接移植这些层会使每个矩阵分块有一半处于空置状态，并在向量运算上形成瓶颈。我们提出了 \textbf{Systolic Linear Attention (SLA)}，围绕矩阵单元重新设计门控 Delta 注意力：利用分组的值头将键的维度扩展至 $256$，将衰减门分成四个各含 $64$ 个通道的分块，并在单个融合内核中完全通过分块矩阵乘积完成分块求解。在 TPU v7x 上，完整的 \texttt{SLA-R4} 层的运行速度是开源 Tokamax KDA 内核的 $1.6\times$--$2.5\times$（采用精确分块求逆时为 $1.4\times$--$2.4\times$），加速主要来自头共享与内核融合。在一项使用 $2{,}048$ 枚 TPU v7x 芯片、涵盖参数量从 $6.5\text{B}$ 到 $63.1\text{B}$ 的稀疏混合专家（MoE）模型的预训练研究中，SLA 在困惑度相近的情况下，训练速度比 KDA 快 $11\%$--$24\%$；在规模较小、随机种子一致的实验中，SLA 达到的困惑度还略低于 KDA。