January 09, 2026
Akshay Mishra, Reiner Pope, Sanjit Neelam, Daniel Heinlein, Vaclav Cvicek, Zaal Vasania, and James Hill-Khurana
Quantizing attention improves efficiency on two fronts: the model has higher compute throughput, and loads fewer bytes per key/value. However, training with block quantized attention can break causal modeling. We present a fix that enables training with MXFP4 in both attention and the attention gradient.
Causal modeling
In causal language modeling, the final logits at position iii must depend only on tokens at positions ≤i\le i≤i. Future leakage is when information from positions >i>i>i may influence the logits at position iii. It poses an issue because it causes a skew between training and decode. In typical setups, causal masks prevent leakage in attention. But block quantization can introduce a subtle new path for future leakage.
Block quantization
Modern accelerators require using block-quantized matrix multiplications for highest throughput. To use these instructions to compute A×BA \times BA×B, the row vectors of AAA and column vectors of BBB must be split into blocks of size kkk and quantized. The specific value of kkk depends on the format being used. For example, the microscaling formats use k=32k=32k=32.
Within a block, let x0,x1,…,xk−1x_0, x_1, \ldots, x_{k-1}x0,x1,…,xk−1 denote the precise (unquantized) elements. When quantizing, we approximate each element xix_ixi with a quantized value qiq_iqi and a scale sss shared across the block: xi≈qi⋅sx_i \approx q_i \cdot sxi≈qi⋅s
There are various approaches to selecting sss with different tradeoffs. But in general they choose an sss with a function of all elements in the block.
Given sss, the quantized elements are:
qi=round(xis)q_i = \text{round}\left(\frac{x_i}{s}\right)qi=round(sxi)
With this procedure, sss and consequently qiq_iqi depend on all the pre-quantized elements in the block. So if the quantization block spans across different token positions, quantization enables the later tokens in the block to influence earlier tokens.
Quantizing attention
Causal attention for a single head takes queries Q\mathbf{Q}Q, keys K\mathbf{K}K, and values V\mathbf{V}V as inputs, and produces:
P=softmax(Q×KT+M)output=P×V\begin{array}{lcl} \mathbf{P} & = & \text{softmax}(\mathbf{Q} \times \mathbf{K}^T + \mathbf{M}) \ \text{output} & = & \mathbf{P} \times \mathbf{V} \end{array}Poutput==softmax(Q×KT+M)P×V
where M\mathbf{M}M is the causal mask.
Our goal is to use block-quantized matrix multiplications for Q×KT\mathbf{Q} \times \mathbf{K}^TQ×KT and P×V\mathbf{P} \times \mathbf{V}P×V. So Q\mathbf{Q}Q and K\mathbf{K}K need to be quantized in blocks formed along the head dimension, while P\mathbf{P}P and V\mathbf{V}V need to be quantized in blocks formed along different token positions.
Quantizing P\mathbf{P}P is safe despite blocking along token positions, since the causal mask zeros out future probabilities. However, the quantized V\mathbf{V}V at position jjj can depend on values at positions >j">j>j, which can cause future leakage.
When does quantized V\mathbf{V}V cause future leakage?
Consider query position iii and value position jjj, with block indices:
bi=⌊ik⌋,bj=⌊jk⌋b_i = \left\lfloor \frac{i}{k} \right\rfloor, \quad b_j = \left\lfloor \frac{j}{k} \right\rfloorbi=⌊ki⌋,bj=⌊kj⌋
Leakage occurs when query position iii and value position jjj share a quantization block position (bi=bjb_i = b_jbi=bj):
- bi=bjb_i = b_jbi=bj (block-diagonal): query iii attends to value jjj if j≤ij \le ij≤i. But the quantized value at jjj is computed from all positions in the block (including positions greater than iii). This breaks causality since the attention output at position iii can depend on positions greater than iii.
- bi>bjb_i > b_jbi>bj (past blocks): All positions in block bjb_jbj precede the first position in block bib_ibi. No leakage.
- bi<bjb_i < b_jbi<bj (future blocks): The causal mask zeros P[i,j]\mathbf{P}\