大语言模型的快速多项式超越

图形处理单元 (GPU) 各代以不同的速率缩放矩阵、特殊功能和内存管道，因此内核瓶颈会随着硬件的发展而变化。 FlashAttention-4 暴露了 NVIDIA Blackwell 上注意力内部的这种不平衡。我们测试短多项式程序是否可以加速大语言模型（LLM）中的其他特殊功能单元（SFU）操作。我们首先将本机 PyTorch 评估与隔离的 IEEE 二进制 16 (FP16) 扫描中的打包融合乘加 (FMA) 程序进行比较，该扫描跨越 L2 驻留和高带宽内存 (HBM) 驻留工作集。然后，我们在四个 GB200 集成任务中用 3 度或 4 度 bfloat16 (BF16) 程序替换原生 sigmoid、tanh 和 sigmoid 线性单元 (SiLU)：密集 SiLU、tanh-softcapped 注意力、sigmoid 注意力和路由专家 Swish 门控线性单元 (SwiGLU)。这些程序将分析对称性、目标格式舍入和消耗内核内的打包算术结合起来。隔离路径在 L2 中提高了 1.19--2.19 倍，在 HBM 中提高了 1.00--1.70 倍。密集 SiLU、tanh-softcapped-attention 和路由专家替换分别将完整训练步骤吞吐量提高了 2.7\%、2.9\% 和 8.0\%。 sigmoid-attention 替换将完全注意力前向提高了 7.4%，完整 GPU 步长提高了 0.3%。每个任务的相同检查点开放权重消融和一对预训练比较将评估扩展到模型行为。在接近 1000 亿个词元的共同范围内，四个任务的最终平滑训练损失差异（多项式减去本机）范围从 $-0.107$ 到 $+0.079$。