Online Softmax: Tiling for Arbitrarily Large Rows

sky_io@outlook.com (K4i) — Tue, 21 Apr 2026 10:00:00 +0800

introduction

in the fused softmax post, we showed that keeping an entire row in GPU SRAM eliminates redundant global memory traffic — reducing softmax from \(8MN\) to \(2MN\) memory operations. the critical assumption: each row of size \(N\) fits in SRAM.

that assumption breaks for large \(N\). modern GPUs have between 48 KB and 228 KB of shared memory per SM. for float32, a row exceeding ~12K–57K elements won’t fit, and the fused approach fails.

Online-Softmax on k4i's blog

Online Softmax: Tiling for Arbitrarily Large Rows

introduction