本文研究能否用短多项式程序加速大语言模型中的特殊函数单元运算。作者在FP16环境下对比了原生PyTorch求值与打包融合乘加程序,并在GB200上以3至4阶BF16多项式替换sigmoid、tanh和SiLU,覆盖稠密SiLU、tanh软上限注意力、sigmoid注意力和路由专家SwiGLU四类任务。隔离测试中,L2驻留场景提速1.19至2.19倍,HBM场景提速1.00至1.70倍;端到端训练步吞吐分别提升2.7%、2.9%和8.0%,sigmoid注意力前向提升7.4%。在约千亿token规模下,四项任务的最终平滑训练损失差异介于-0.107至+0.079之间,表明该替换在提速的同时基本不损害模型行为。
| SFU | 特殊函数单元,GPU 中用于执行超越函数(如 sigmoid、tanh)等复杂数学运算的硬件组件。 |
| FMA | 融合乘加,一种将乘法和加法合并为单条指令的运算,可提高计算效率和精度。 |
| BF16 | bfloat16 浮点格式,一种 16 位浮点数表示,具有较宽的动态范围,常用于深度学习训练。 |
| SiLU | Sigmoid 线性单元,一种激活函数,定义为 x * sigmoid(x),也称为 Swish。 |
| SwiGLU | Swish 门控线性单元,一种结合 Swish 激活和门控机制的神经网络层,常用于 Transformer 的前馈网络。 |
📱 每天一份 AI 前沿日报
关注公众号,每天 09:00 推送 · 不错过任何重磅