Future leakage in block-quantized attention | MatX

Quantizing attention improves efficiency on two fronts: the model has higher compute throughput, and loads fewer bytes per key/value. However, training with blo

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 ii must depend only on tokens at positions ≤i\le i. Future leakage is when…

saved by

• Aaron Pham
• Akira Yoshiyama

related reading

• [2507.04239] Scaling Context Requires Rethinking Attentionarxiv.org
• Linear Attention Fundamentals | Hailey Schoelkopfhaileyschoelkopf.github.io
• [2205.14135] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenessarxiv.org
• 2502.11089arxiv.org
• Log-Linear Attentionarxiv.org
• Overleaf Examplearxiv.org
• Linear Transformers Are Faster After All – Manifest AImanifestai.com
• 1.5x faster MoE training with custom MXFP8 kernels · Cursorcursor.com
• Kimi Linear: An Expressive, Efficient Attention Architecturearxiv.org
• Sparse Attention Post-Training for Mechanistic Interpretabilityarxiv.org
• Scalable MatMul-free Language Modelingarxiv.org
• A User's Guide to FlexAttention in FlashAttention CuTe DSL - Colfax Researchresearch.colfax-intl.com