Where in attention is FP32 still used, and what breaks if you push everything to FP16/BF16?
Matmuls run in FP16/BF16 to use tensor cores; softmax stays in FP32 (or uses running-max stabilization) because exp overflows FP16's ~65504 ceiling.
Imagine a tiny notebook that can only write numbers up to about 65,000. Most of the math in attention is fine, because you're adding lots of small numbers and the result stays in range. But one step asks you to write down e raised to a fairly big number. e to the 11 is already around 60,000. e to the 12 doesn't fit at all. So that one step either needs a bigger notebook with more room, or a clever trick where you shift every number down by the largest one before you write it. Picture using a regular pad for most work and pulling out a giant whiteboard only for that one tricky step.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Mixed-precision attention is one of those engineering details where a sloppy implementation produces a model that trains, looks fine on the loss curve, and underperforms its proper-precision twin by several points on downstream evals. The bug is silent. The fix is mechanical once you see the structure.
The shape worth memorizing: matmuls run in low precision because tensor cores deliver peak throughput there and sum of products is forgiving of per element rounding. The softmax stays in higher precision (or uses an online-softmax stabilization) because exp is not forgiving. It amplifies its input exponentially, and FP16 simply does not have the dynamic range to represent the result.
We will walk why matmuls forgive, why softmax does not, how BF16 and FP8 reshape the picture, and how FlashAttention bakes the whole stability story into a single kernel, all numbers and recipes current to 2026 hardware (Hopper / Blackwell).
Why matmuls are fine in low precision
Matmul is structurally a sum of many products. The reduction has a useful averaging property: per element rounding errors do not add coherently, they cancel partially across the sum. Even when individual products are slightly off, the final accumulated value is close to what FP32 would have produced.
Tensor cores exploit this. They hit peak throughput in FP16, BF16, and FP8, and they accumulate in FP32 internally regardless of input precision. Hopper's HMMA instructions and Blackwell's TMMA do this natively. You get the throughput of the low precision input format with the precision of an FP32 accumulator on the output side.
In attention specifically, this covers a lot of ground. Q, K, V projections are matmuls. QK^T is a matmul. The product against V is a matmul. The output projection is a matmul. All run safely in FP16/BF16 with FP32 accumulators. The per element error magnitude in a single FP16 product is around 1e-3, but summing 128 such products gives final error around 1e-3 × √128 ≈ 0.011, well within tolerance for the next layer.
Matmuls are the part of the attention pipeline you actively want in low precision, that is where tensor core utilization is highest.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
| Operation | Typical precision | Why |
|---|---|---|
| Q,K,V projections | FP16/BF16 (FP8 on H100) | Tensor core throughput; robust to element-wise error |
| QKᵀ matmul | FP16/BF16 inputs, FP32 accumulator | Tensor core throughput; accumulator preserves precision |
| Softmax exp + sum | FP32 (or stabilized online softmax) | Avoid exp overflow; preserve small probabilities |
| @V matmul | FP16/BF16 inputs, FP32 accumulator | Same as QKᵀ |
| Output projection | FP16/BF16 | Standard low precision matmul |
Real products, models, and research that use this idea.
- PyTorch scaled_dot_product_attention casts softmax intermediates to FP32 for stability across backends.
- FlashAttention 2 and 3 maintain FP32 running max and denominator across SRAM tiles inside the kernel.
What an interviewer would ask next. Try answering before peeking at the approach.
QWhy doesn't BF16 fully replace FP32 for softmax stats?
BF16 has FP32's exponent range (no overflow) but only 7 mantissa bits. The softmax denominator can be a sum of many small exp values; in BF16 those small values lose mantissa precision and round to the same quantized levels, degrading the running denominator. FP32 accumulators preserve precision in the sum.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Thinking the matmuls are the problem and softmax is fine. It's the opposite: matmuls are robust in FP16; softmax's exponential is what overflows.
60 second bullets to scan on the way to the call.
Which ops run in FP16/BF16 and which need FP32
Why FP16 overflows for exp (max ~65504, exp(11) already there)
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.