面向 LLMs 的快速多项式超越函数

arXiv:2610.00049v1 发布类型：新 摘要：GPU 各代产品以不同速率扩展矩阵、特殊函数和内存流水线，因此随着硬件演进，内核瓶颈也会转移。FlashAttention-4 暴露了 NVIDIA Blackwell 注意力计算中的这种不平衡。我们测试短小的多项式程序能否加速大语言模型（LLMs）中的其他特殊函数单元（SFU）操作。首先，我们在独立的 IEEE binary16（FP16）测试中，将原生 PyTorch 计算与打包融合乘加（FMA）程序进行比较；该测试覆盖驻留于 L2 和高带宽内存（HBM）中的工作集。随后，在四项 GB200 集成任务中，我们用三阶或四阶 bfloat16（BF16）程序替换原生 sigmoid、tanh 和 sigmoid 线性单元（SiLU）：稠密 SiLU、tanh 软上限注意力、sigmoid 注意力，以及路由专家 Swish 门控线性单元（SwiGLU）。这些程序在使用它们的内核内部结合了解析对称性、目标格式舍入和打包算术。独立路径在 L2 中的性能提升为 1.19–2.19 倍，在 HBM 中为 1.00–1.70 倍。稠密 SiLU、tanh 软上限注意力和路由专家替换分别使完整训练步骤吞吐量提升 2.7%、2.9% 和 8.0%。sigmoid 注意力替换使完整注意力前向计算性能提升 7.4%，使完整 GPU 步骤性能提升 0.3%。采用相同检查点的开放权重消融实验，以及每项任务一次的配对预训练比较，将评估扩展到了模型行为。在接近 1000 亿 tokens 的常见训练时长处，四项任务中最终平滑训练损失的差异（多项式减原生）介于 $-0.107$ 和 $+0.079$ 之间。