Deep Learning Fundamentals

What is a fused softmax kernel, and why is it faster?

By Interview Clozo · Published 27 October 2026 · Last verified 26 October 2026 · 13 min read

A fused softmax kernel is one graphics processing unit (GPU) program that runs every step of softmax (maximum, subtract, exponentiate, sum, divide) on a row held in fast on-chip memory. Softmax does little arithmetic per byte it moves, so fusing saves time by cutting memory traffic, not arithmetic. Counting reads and writes predicts about four times less memory traffic,¹ but that count compares against unfused operations, which is why a real benchmark can show no gain.

In this article

Before you start
Why interviewers ask this
Fused softmax kernel from first principles
Worked example
Trade-offs and edge cases
How to answer it in an interview
Common mistakes
Follow-up questions
Sources

Before you start

This page is one step of the neural network learning path.

Why interviewers ask this

The question separates candidates who know what softmax computes from candidates who know what a GPU spends time on. Softmax is cheap arithmetic, so a design discussion about attention or a custom kernel soon turns to memory traffic.

A weak answer says "fusing makes it faster because there is one function call instead of many." A strong answer says that the operation is limited by memory bandwidth, counts the reads and writes with and without fusion, and states the limit: the row must fit on chip. A Staff-level answer adds online softmax, which removes the extra pass for the maximum and lets long rows be processed in tiles, and notes that the theoretical gain is not what a benchmark shows at every size.

Fused softmax kernel from first principles

Fused softmax kernel: A fused softmax kernel is a single GPU program that computes every step of softmax for a row (maximum, subtraction, exponentiation, sum and division) while the row stays in fast on-chip memory, so each input value is read from main memory once and each output value is written once.
Safe softmax: Safe softmax is a form of softmax that is computed after subtracting the row's largest value from every score. The shift leaves the result unchanged but keeps every exponent at or below zero, so the exponentials cannot overflow. It costs one extra pass over the row to find the maximum first.

In plain words

Picture a cook who needs five steps to prepare a dish. In the slow kitchen the cook walks to the pantry before every step and walks back to put the result away. In the fast kitchen the cook carries the ingredients to the counter once, does all five steps there, and walks back once. The chopping is the same; only the walking changes.

The analogy breaks in one place. The counter is small. If the ingredients do not fit, the cook must work in portions, which is what the last step of this page covers.

A kernel is a small program that runs on the GPU across many threads. The main idea in one precise sentence: a fused softmax kernel does the same arithmetic as separate softmax operations but moves each number between DRAM and the chip about once, which matters because the time goes into the moves.

Step 1: Safe softmax reads each row three times

Safe softmax reads each row three times: once to find the maximum, once to sum the exponentials, and once to write the outputs. Softmax is $y_{i} = e^{x_{i}} / \sum_{j} e^{x_{j}}$ over the scores of one row.

Subtracting any constant $c$ from every score multiplies the top and the bottom of the fraction by the same factor, so the result does not change. Choosing $c$ as the row maximum keeps every exponent at or below zero. Without it, the exponential of a large score can overflow, and the exponential of a very negative one can underflow to zero. Milakov and Gimelshein (2018) note that all major deep learning frameworks used this safe form, and that it does three passes over the input, giving four memory accesses per element overall.

Input: one row of scores. Output: one row of weights. You now know the work; the next step counts what moving it costs.

Step 2: Fusing keeps the row on chip and cuts memory traffic

A fused kernel loads each row once, does every step in on-chip memory, and writes the result once. The Triton fused softmax tutorial counts the alternative for a matrix of M rows and N columns: separate operations read 5·M·N + 2·M elements from DRAM and write back 3·M·N + 2·M elements. A fused kernel reads and writes only M·N elements each. The tutorial gives the theoretical speed-up as about four times, the ratio (8·M·N + 4·M) / (2·M·N).

The tutorial's code comments show the fused flow: load the row into SRAM, subtract the maximum for numerical stability, then write the output back to DRAM. It also says the exponential in Triton is fast but approximate, like __expf.

Input: a row in DRAM. Output: the weights in DRAM, with no intermediate arrays. The four times is a bound on memory traffic against separate operations, not a promise about time. A built-in that is already optimised moves less than the separate operations do, so the measured speed-up can be far smaller, and the worked example shows widths where it is negative.

Step 3: Online softmax finds the maximum and the sum in one pass

Online softmax computes the row maximum and the sum of exponentials together in one pass, rescaling a running sum whenever a new maximum appears. For each new score $x_{j}$ the update from Milakov and Gimelshein (2018) is $m_{j} = \max (m_{j - 1}, x_{j})$ and $d_{j} = d_{j - 1} e^{m_{j - 1} - m_{j}} + e^{x_{j} - m_{j}}$ . Here $m$ is the running maximum and $d$ the running sum.

The rescaling works because the old sum was measured against the old maximum. Multiplying by $e^{m_{j - 1} - m_{j}}$ re-expresses it against the new maximum, and the paper proves the result by induction. The outputs are then $e^{x_{i} - m} / d$ for the final $m$ and $d$ .

Compared with an unfused safe softmax that re-reads the row from memory on each pass, this cuts accesses from four to three per element. Inside a fused kernel whose row is already on chip, it saves an on-chip pass and a reduction, and its memory saving matters once the row is too long to fit. The combine operation is associative, so partial results from different blocks of a row can be merged in any grouping. That property is what makes the next step possible.

Step 4: Tiling extends the idea to rows that do not fit

Tiling splits a long row into blocks and merges the per-block results with the online rescaling, so the row no longer has to fit in on-chip memory. The Triton tutorial's kernel assumes the row fits in SRAM, and because Triton blocks must have a power-of-two number of elements, it pads each row internally.

FlashAttention applies the tiled form to attention. Dao et al. (2022) split the matrices Q, K and V (query, key and value) into blocks, load them from slow high-bandwidth memory (HBM), the GPU's DRAM, into fast SRAM, and scale each block's output by the right normalization factor so the final result is correct. The large sequence-by-sequence score matrix is never written to HBM. A standalone tiled softmax still reads the row twice, once for the maximum and sum and once to write normalised outputs; FlashAttention needs one pass only because it also fuses the multiplication by V. FlashAttention-2 states that online softmax can split attention into blocks and rescale each block's output to get the right result with no approximation.

You now have the whole mechanism: three passes, one load and one store, one pass for the sum, and blocks for long rows.

Fusing removes memory trips, online softmax removes the extra pass for the maximum, and tiles remove the row-length limit.CC BY 4.0 — source: What is a fused softmax kernel, and why is it faster?

Step-by-step description

Three passes per row.
One load, one store.
Maximum and sum in one pass.
Blocks for long rows.

Safe softmax reads each row three times.
A fused kernel loads each row once and writes the result once.
Online softmax finds the maximum and the sum in one pass.
Tiling lets rows that do not fit on chip be processed in blocks.

Knowing it is half the job. Saying it is the other half.

Interview Clozo is an AI interview copilot for real interview rounds: it hears the question and puts a strong answer on your screen, and you say it in your own words. See the features or how it works.

Start Free Trial →

Worked example

Take an illustrative matrix of 4,096 rows and 1,024 columns in 32-bit floating point (fp32), four bytes per number.¹ One pass over it moves 4,194,304 elements, or 16.8 megabytes. Using the Triton formulas, the separate operations move 20,979,712 elements in and 12,591,104 out, 33,570,816 in total, or 134.3 megabytes. The fused kernel moves 8,388,608 elements, or 33.6 megabytes, so the ratio is 4.002.¹

For this matrix, separate softmax operations move about 134.3 megabytes against 33.6 for the fused kernel, a ratio of 4.002.
Quantity	Value (computed, illustrative matrix)
Elements in one pass over the matrix	4,194,304
Megabytes in one pass	16.8
Elements moved, separate operations	33,570,816
Megabytes moved, separate operations	134.3
Elements moved, fused kernel	8,388,608
Megabytes moved, fused kernel	33.6
Traffic ratio	4.002

Fusing the row cuts the memory traffic of the same arithmetic to about a quarter.CC BY 4.0 — source: What is a fused softmax kernel, and why is it faster?

Text description of the diagram

Separate operations move 134.3 megabytes.
The fused kernel moves 33.6 megabytes.

Now the online recurrence on three scores, 2, 5 and 3, all illustrative.¹ Start with $m = 2$ and $d = 1$ . The score 5 raises the maximum, so $d = 1 \cdot e^{- 3} + 1 = 1.0498$ and $m = 5$ . The score 3 does not, so $d = 1.0498 + e^{- 2} = 1.1851$ . The outputs are 0.0420, 0.8438 and 0.1142, which sum to one and match safe softmax.¹

The running sum changes only when a score arrives, and the final outputs equal safe softmax.
After score	Running maximum	Running sum
2	2	1.0000
5	5	1.0498
3	5	1.1851

The safe shift matters for large scores. For the scores 1000 and 1001, the naive exponential overflows in double precision, while the safe form gives 0.2689 and 0.7311.¹ This is overflow in the forward calculation, a different problem from exploding gradients in training.

Measured speed differs from the bound. In the Triton tutorial's benchmark of one GPU run (the page does not name the hardware), the matrix is fixed at 4,096 rows in fp32 and the number of columns varies. At 4,096 columns the Triton kernel reaches 1404.7 gigabytes per second, the built-in torch.softmax 1314.0 and a naive version 338.6.¹ That is 1.07 times the built-in and 4.15 times the naive one.

At 2,176 columns the figures are 1253.6, 960.7 and 325.7, a ratio of 1.30 over the built-in. In this run the Triton kernel was behind torch.softmax at every width listed from 256 to 1,024 columns: 511.2 against 689.7 gigabytes per second at 256 columns, a ratio of 0.74, 831.2 against 914.7 at 512, 907.1 against 973.6 at 768, and 1035.0 against 1068.9 at 1,024, a ratio of 0.97.¹ The four-times bound fits the comparison with the naive version, not with an already optimised built-in.

Trade-offs and edge cases

Fusion trades generality for traffic. The Triton kernel works for rows that fit in SRAM and relies on padding and masking to reach a power-of-two block size. Longer rows need the tiled online form of Step 4, which does a few extra operations per element; Milakov and Gimelshein call the added cost negligible.

Online softmax alone is a modest gain. Milakov and Gimelshein report up to 1.3 times over safe softmax (paper abstract), measured at a vector length of 4,000 with a batch of 4,000 vectors on a Tesla V100 in fp32. That is close to the 1.33 times that comes from going from four accesses to three against an unfused safe softmax.

The larger gains come from fusing with a neighbouring operation, such as top-k, the selection of the k largest values. Run separately, safe softmax and top-k need five accesses per element, and fused they need one. The paper's measured speed-up there is up to five times at small k, which it splits as 2.5 times from fusion and 2 times from online softmax.

The fused kernel needs the row to fit on chip, and the online tiled form removes that limit at a small arithmetic cost.

The fused kernel needs the row to fit on chip; the online tiled form removes that limit.CC BY 4.0 — source: What is a fused softmax kernel, and why is it faster?

Text description of the diagram

A fused kernel needs the row to fit in SRAM.
An online tiled kernel works for long rows at a small extra arithmetic cost.

Edge cases: the fast exponential is approximate, so a fused kernel can differ from the built-in in the last digits; lower-precision number formats add their own care, covered in half-precision formats; and a fused kernel that wins on one GPU and row size can lose on another, as the benchmark rows show. Rely on measurement for the real decision.