arXiv:cs.LG(机器学习,全量分类)· Robert Hu·· 14 小时前AI 评分35
用快速多项式超越函数加速 LLM 中的 SFU 运算
Fast Polynomial Transcendentals for LLMs
AI 导读
研究测试短多项式程序能否加速 LLM 中的特殊函数单元(SFU)运算,在 GB200 上用 3 次或 4 次 BF16 多项式替换原生 sigmoid、tanh 和 SiLU。
正文
Abstract:Graphics processing unit (GPU) generations scale matrix, special-function, and memory pipelines at different rates, so kernel bottlenecks move as hardware evolves. FlashAttention-4 exposed this imbalance inside attention on NVIDIA Blackwell. We test whether short polynomial programs can accelerate other special-function-unit (SFU) operations in large language models (LLMs). We first compare native PyTorch evaluation with packed fused multiply--add (FMA) programs in an isolated IEEE binary16 (FP16) sweep spanning L2-resident and high-bandwidth-memory (HBM)-resident working sets. We then replace native sigmoid, tanh, and sigmoid linear unit (SiLU) with degree-3 or degree-4 bfloat16 (BF16) programs in four GB200 integration tasks: dense SiLU, tanh-softcapped attention, sigmoid attention, and routed-expert Swish-gated linear unit (SwiGLU). The programs combine analytical symmetry, target-format rounding, and packed arithmetic inside consuming kernels. The isolated paths improve by 1.19--2.19x in L2 and 1.00--1.70x in HBM. The dense-SiLU, tanh-softcapped-attention, and routed-expert substitutions improve complete training-step throughput by 2.7\%, 2.9\%, and 8.0\%, respectively. The sigmoid-attention substitution improves complete-attention forward by 7.4\% and the complete GPU step by 0.3\%. Same-checkpoint open-weight ablations and one paired pre-training comparison per task extend the evaluation to model behavior. At common horizons near 100 billion tokens, the final smoothed training-loss differences (polynomial minus native) range from $-0.107$ to $+0.079$ across the four tasks.
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2610.00049 [cs.LG] |
| (or arXiv:2610.00049v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.00049 arXiv-issued DOI via DataCite |
Submission history
From: Robert Hu [view email]
[v1]
Thu, 3 Sep 2026 17:18:10 UTC (404 KB)
来源:arXiv:cs.LG(机器学习,全量分类) · arxiv.org