cs.LGSep 3, 2026

Fast Polynomial Transcendentals for LLMs

Authors: Robert Hu

Organizations: GRAPHCORE RESEARCH

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-0.107 to +0.079+0.079 across the four tasks.

Figures & tables

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

Sep 3, 2026

Hardware-Aware FP4 FlashAttention-4

Blackwell's 4-bit floating-point (FP4) tensor cores do not automatically make attention faster because softmax conversion and on-chip dependencies dominate once its matrix products shrink. We address this with Direct-P for noncausal inference and a causal path that passes the forward quantization directly into backward. Direct-P maps scores directly to FP4 probabilities and reaches up to 2.13×\times the bfloat16 (BF16) forward throughput on an NVIDIA GB200. The causal path reconstructs probabilities from saved quantized queries and keys and uses 8-bit floating-point (FP8) gradient operands, accelerating a complete single-GPU 8-billion-parameter update by up to 1.14×\times. Matched distributed training retains FP8 probabilities and values; every tested MXFP4 probability/value training trajectory diverges.
Sep 30, 2026

LampAttention: Look-Ahead Mixed-Precision FlashAttention for Dedicated Accelerators

While most attention logits can be computed in low precision without degrading numerical stability, current attention kernels fail to exploit this phenomenon. We introduce a novel hardware-algorithm co-design in the form of mixed-precision FlashAttention. Our method accumulates key-query products and evaluates their exponentials in 8-bit formats, then adaptively identifies sensitive sub-blocks and recomputes them in 16-bit formats. We propose the specifications for a dedicated accelerator capable of executing this pipeline efficiently. Simulated experiments with Qwen3 and Gemma 3 show that rerouting a selective minority of sub-blocks to high precision is sufficient to recover the baseline model performance.
May 5, 2026

Tile-Level Activation Overlap for Efficient LLM Inference

SwiGLU is the dominant MLP activation in modern large language models, yet its intermediate tensor materialization costs 9-37% of MLP execution time. We present two complementary CUTLASS-based SM90 kernels that fuse SwiGLU into GeMM at the tile level. Kernel-1 overlaps Swish computation on the Gate accumulator with Up-tile loading using the Pingpong warp-specialized schedule; Kernel-2 interleaves SwiGLU with tile stores via a custom Epilogue Visitor Tree. Evaluated on Qwen-2.5 models (0.5B-72B) on NVIDIA H100, our kernels achieve up to 2.47x speedup over PyTorch, shifting workloads from memory-bound to compute-bound and reaching 79.5% peak BF16 utilization. We demonstrate that torch.compile cannot replicate this fusion (3-7x slower than our kernels), validating the need for hand-crafted tile-level design. Our fused kernels are also numerically superior, achieving zero mismatches compared to 4.5-11% for cuBLAS.