💻 computer science

P-Cast Precision in FP8 Attention: Sink-Induced Collapse and the Optimality of S=2^8

This paper analyzes FP8 attention precision loss caused by the Attention Sink phenomenon, demonstrating that reverse KV iteration and a static scaling factor of $S=256$ effectively prevent probability underflow and minimize quantization error, thereby explaining the engineering choices in FlashAttention-3/4 and providing a closed-form threshold for predicting precision loss.

Original authors: Reed Lau

Published 2026-06-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Reed Lau

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Tiny Bucket" Problem

Imagine you are a chef trying to pour a massive amount of soup (data) into a very small, specific-sized cup (the computer's memory format called FP8).

In modern AI, we want to cook faster, so we use these tiny cups. However, this cup has a flaw: it can only hold a few specific sizes of liquid. If the soup is too thin (too small a number), the cup treats it as "nothing" and pours it out. If the soup is too thick, it overflows.

The paper focuses on a specific problem in AI called Attention. Think of Attention as the AI deciding which words in a sentence are important. Usually, the AI pays attention to the most recent words. But sometimes, it gets obsessed with the very first words of a sentence (called "Sink Tokens"). These first words become so loud and important that they drown out everything else.

When the AI tries to squeeze these "loud" first words and the "quiet" later words into our tiny FP8 cup, the quiet words get crushed. They become so small that the cup sees them as zero. This is called P-Collapse. The AI forgets the middle of the sentence entirely.

The Two Fixes: Changing the Order and the Scale

The authors found two simple ways to fix this "crushing" problem without changing the hardware.

Fix 1: The "Reverse Order" Strategy

The Problem: Imagine you are filling a bucket with water. If you pour a giant firehose (the loud first word) in first, the water level rises instantly. When you try to add a single drop (the quiet later words) afterward, the drop is so small compared to the firehose that it disappears into the background noise.

The Solution: Instead of pouring the firehose first, pour the single drops first, and save the firehose for last.

How it works: The AI processes the sentence from the end to the beginning (Reverse Iteration).
The Result: The quiet words are processed while the "water level" is still low, so they fit perfectly in the cup. The loud firehose (the first word) is added at the very end, where it doesn't crush the previous drops.

Fix 2: The "Magnifying Glass" Strategy (The $S=256$ Discovery)

The Problem: Even if you pour the drops carefully, the cup is still too coarse. It has "steps" or "rungs" on a ladder. If a drop falls between two rungs, it gets rounded down to the lower rung. If it's too small, it falls off the bottom rung and becomes zero.

The Solution: Before pouring the soup into the cup, the authors suggest using a magnifying glass (a scaling factor) to make the soup look bigger.

The Magic Number: They tested many magnification levels. They found that 256 is the perfect number.
Why 256?
1. It's a "Clean" Number: In computer math, 256 is a power of two ( $2^8$ ). This means the computer can multiply and divide by 256 perfectly without losing any tiny bits of information (unlike other numbers like 448, which introduce tiny errors).
2. It Fits the Ladder: The FP8 cup has a specific shape. 256 aligns perfectly with the "rungs" of the ladder, ensuring the smallest drops don't fall off the bottom.
3. It's the Biggest Safe Size: You can't use a magnifying glass that is too strong (like 512), or the big firehose will overflow the cup. 256 is the strongest magnifying glass that keeps everything inside the cup without spilling.

The "Sawtooth" Discovery

The authors drew a graph that looks like a sawtooth (a jagged mountain range).

The "valleys" of the sawtooth represent the best possible precision.
They discovered that only the powers of two (1, 2, 4, 8... 256) sit in these perfect valleys.
Other numbers, like 448 (which some other AI systems use), sit on the "slopes" of the sawtooth. They are okay, but they are slightly less precise than 256.

The Results: What Happened?

The researchers ran tests to see what happens when the AI gets obsessed with the first word (a "Sink Strength" of about 7).

Before the fix (using the standard method): The AI lost 30% to 40% of the important information from the middle of the sentence. It was like trying to read a book where half the pages were blank.
After the fix (using Reverse Order or S=256): The error dropped by 3 to 10 times. The AI could "hear" the quiet words again.
The Best Combo: Using both fixes together didn't make it much better than using just one. It's like having a seatbelt and an airbag; once you have the seatbelt (one fix), the airbag (the other fix) doesn't add much extra safety because the first one already solved the main problem.

The Bottom Line for Engineers

The paper concludes that if you are building an AI system that uses this specific "tiny cup" (FP8 E4M3):

Stop using the number 1 (direct casting) or 448 as your scaling factor.
Switch to 256. It is the mathematically perfect "sweet spot" that keeps the computer's math clean and prevents the AI from forgetting the middle of the sentence.
Or, process the sentence backwards. This also solves the problem, but changing the number to 256 is often easier to implement.

The authors have already updated their own software (hpc-ops) to use 256, and they suggest others do the same to get a free, instant improvement in accuracy.

Technical Summary: P-Cast Precision in FP8 Attention

1. Problem Statement

The paper addresses a critical precision bottleneck in FP8 (E4M3) attention acceleration for Large Language Models (LLMs). While casting the softmax probability matrix $P$ to E4M3 before the $P \cdot V$ matrix multiplication offers significant throughput gains, the format's 3-bit mantissa introduces severe quantization errors under the Attention Sink phenomenon.

In the Attention Sink phenomenon, initial tokens receive disproportionately large attention weights. When combined with forward KV block iteration (processing tokens $0 \to N$ ), the high logit scores of sink tokens inflate the running softmax maximum ( $m$ ). This forces subsequent non-sink probability values ( $P$ ) to become extremely small. Under E4M3 quantization, these small values fall below the representable range and underflow to zero (a phenomenon termed P-collapse).

The paper identifies two implementation choices that determine whether this collapse occurs:

KV Block Iteration Order: Whether the online softmax processes blocks forward ( $0 \to N$ ) or reverse ( $N \to 0$ ).
Static Scaling Factor ( $S$ ): The constant multiplier applied to $P$ before casting to E4M3 (and divided back in the epilogue).

2. Methodology and Analysis

Theoretical Framework

The authors derive a closed-form expression for the fraction of non-sink $P$ values that underflow to zero under forward iteration with a scale factor $S$ :
$F(\Delta, S) = \Phi(\Delta + \delta_k - 6.93 - \ln S)$
Where:

$\Delta$ is the sink strength (logit gap between sink and non-sink tokens).
$\delta_k$ is the expected maximum of $k$ standard Gaussian variables within the sink block (approx. 1.0 for $k=4$ ).
$6.93 \approx 10 \ln 2$ represents the E4M3 round-to-zero threshold ( $2^{-10}$ ).
$\Phi$ is the standard normal CDF.

This analysis reveals that P-collapse is a threshold effect activating around $\Delta \approx 6\text{--}7$ , where a significant fraction of non-sink probability mass is silently zeroed, leading to substantial Mean Squared Error (MSE) in the output.

Optimization 1: Reverse KV Iteration

The paper proves that reversing the iteration order ( $N \to 0$ ) defers the processing of the sink block to the final step. During all preceding iterations, the running maximum is determined only by non-sink tokens. Theoretical bounds show that for practical sequence lengths ( $N \le 10^6$ ), the pre-sink maximum remains low enough that, when combined with $S=256$ , the probability of any non-sink $P$ value underflowing is effectively zero ( $< 10^{-12}$ ).

Optimization 2: The Optimality of $S = 256$

The paper provides a constructive characterization of why $S=256$ ( $2^8$ ) is the optimal static scale, satisfying three simultaneous conditions:

Bit-Exact Scaling: As a power of two ( $2^k$ ), multiplication and division by $S$ are exact in IEEE 754 FP32, introducing no rounding error during the scaling/unscaling process. Non-power-of-two scales (e.g., $S=448$ ) introduce systematic rounding errors.
Minimum Quantization Step ( $d_p$ ): The authors define $d_p(S)$ as the normalized maximum quantization step. They demonstrate that $d_p(S)$ forms a "sawtooth" function over the E4M3 number line, where power-of-two values ( $S=2^k$ ) sit on the lower envelope with a minimum value of $2^{-4}$ . $S=448$ sits 14% higher on this sawtooth, implying a larger worst-case quantization error.
Maximum Normal-Range Coverage: Among power-of-two scales that do not cause overflow ( $S \le 448$ ), $S=256$ is the largest, providing the deepest coverage into the subnormal region of the $P$ domain ( $2^{-6}/256 = 2^{-14}$ ).

3. Key Contributions

Quantification of P-Collapse: The paper derives the closed-form threshold $\Delta_c = 6.93 + \ln S - \delta_k$ for predicting when P-cast failure occurs, moving beyond empirical observation to analytical prediction.
Proof of Sufficiency: It rigorously demonstrates that reverse iteration combined with $S=256$ guarantees zero underflow for non-sink values in practical scenarios.
Characterization of $S=256$ : It establishes $S=256$ as the unique scale satisfying bit-exactness, minimum quantization step, and maximum normal coverage among $2^k$ scales.
Kernel-Faithful Validation: The authors simulate production kernel semantics (isolating the P-cast effect by keeping Q, K, V in FP32) to validate their theoretical claims.

4. Experimental Results

Using kernel-faithful simulations with $N=4096$ and varying sink strengths ( $\Delta$ ):

MSE Improvement: At moderate sink strengths ( $\Delta \approx 7$ ), the baseline configuration (Forward iteration, $S=1$ ) exhibits 3–10× higher MSE compared to optimized configurations.
Comparison of Scales: Forward iteration with $S=256$ yields approximately 10–15% lower MSE than Forward iteration with $S=448$ (used in FlashInfer and TensorRT-LLM), consistent with the theoretical $d_p$ ratio.
Saturation: When both optimizations (Reverse + $S=256$ ) are applied, the results are statistically indistinguishable from Forward + $S=256$ in the P-collapse regime. Both reach the same precision floor determined by inherent E4M3 quantization noise.
Sequence Length: The MSE improvement of optimized configurations over the baseline grows with sequence length, as longer sequences increase the non-sink probability mass at risk.

5. Significance and Claims

The paper positions itself not as a proposal of new techniques, but as a quantitative explanation for implementation choices already adopted in FlashAttention-3/4 (which use Reverse + $S=256$ ).

Practical Impact: The analysis motivated the update of the hpc-ops kernel from $S=1$ to $S=256$ .
Recommendations: The authors recommend that implementations currently using $S=1$ or $S=448$ switch to $S=256$ . This is a "zero-cost" change (bit-exact scaling adds negligible latency) that removes P-collapse and strictly improves or maintains P-cast precision.
Modesty: The authors explicitly state that while the P-cast MSE improvement is significant (3–10×), the end-to-end impact on model perplexity or task accuracy is not guaranteed, as it may be dominated by other FP8 noise sources (QKV quantization). They frame the switch to $S=256$ as "strictly-no-worse" rather than a guaranteed quality win for all downstream metrics.
Scope: The analysis is specific to E4M3 P-casting and does not claim transferability to other formats (e.g., MXFP4) or per-channel quantization strategies where preconditions differ.