FLAMIA: FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell
Fedor Velikonivtsev, Mikhail Koltakov, Denis Kuznedelev
cs.LG
2026-11-03
Yandex 的 Blackwell 低比特注意力训练核:前反向合计提速 1.06–1.42 倍,1B 预训练损失几乎不变,量化尺度的因果泄露也堵上了。
Blackwell 的 tensor core 上,FP8 峰值吞吐是 BF16 的两倍、FP4 四倍,FlashAttention-4 却停在 BF16。低精度用进注意力训练卡在三件事上。一是量化开销:注意力的六个操作数各自参与两次收缩方向不同的矩阵乘,常规块缩放要为每个操作数存两份量化表示,而量化是访存受限的,省下的算力会被它吃掉。二是因果泄露:causal mask 挡不住量化尺度带来的依赖,未来 token 抬高共享尺度的最大值,会顺带改变更早 token 的输出。三是前反向不一致:两边在量化、分块或指数求值上对不齐,单步梯度相似度再高,训练照样发散。现有低比特实现各有缺口:FlashInfer 只有 packed 前向,cuDNN FP8 用整张量尺度,Low Precision FA4 不支持 causal mask。
FLAMIA 有 MXFP8 与 NVFP4/FP8 混合两条训练路径,在 FlashAttention-4 的 warp-specialized 结构上扩展。
核心设计是 token-block scaling:32 个连续 token 共享一个 2 的幂指数,覆盖整个 head 维度,这一块的编码和尺度在转置后依然有效;8-bit MMA 又同时吃 K-major 和 MN-major 两种布局,一份 E4M3 载荷就能喂两个方向的矩阵乘,重复量化就此省掉。NVFP4 只用于 QK^T 及其反向重算,两者共享收缩轴,一份表示两用,其余乘法走 8-bit。
因果安全的处理很克制:尺度组至多覆盖 32 个连续位置;尺度指数向上取整,元素永不饱和;4-bit 的 V 只在跨因果边界的对角 tile 改用每 key 独立尺度并折进 P̂,其余 tile 保留便宜的分组尺度,4096 token 时前向代价从 2.33× 降到 1.17×。反向的 FP16 dQ 累加器默认用整序列统计定标,causal dQ scale 选项可以把依赖收回到行内。
前反向一致性靠复刻:反向用与正向相同的量化表示、tiling、对角决策和指数求值重算概率;dS 量化成 E4M3,尺度取 32×32 块内最大绝对值,一份载荷同时服务 dQ 和 dK。
packed 文档原布局直过:量化统计不跨文档边界,scale 表按 token 密集存储,调度在 GPU 上按 tile 代价排序,131072 个单 token 文档的工作列表构建 0.26 ms,参考实现要 50.5 ms;T=131072 时峰值显存 2.7 GB,对齐 padding 方案要 16.8 GB。全 4-bit 路径存在但非默认:4-bit MMA 没有 MN-major 形式,六份量化拷贝占反向时间 45–52%,更慢且 loss 更高。位级指数近似(Mitchell-P、Schraudolph)只留给推理,训练不用。
| 配置 | 指标 | 结果 |
| MXFP8,B200,9 个 causal packed 负载 | 前反向+量化提速,对每负载最快基线 | 1.10–1.40× |
| MXFP8,B300 | 同上 | 1.06–1.42× |
| NVFP4,kernel-only causal 反向 | 提速对最快基线 | B200 1.22–2.01×,B300 1.17–2.00× |
| MXFP8 数值 | O/dQ cosine 对 FP32 | 0.9988/0.9979(cuDNN MXFP8 为 0.9988/0.9978) |
| 1B Llama-3,100B tokens,8×B300 | 训练 loss | BF16 2.392,MXFP8 2.397,NVFP4 2.406 |
| 同上 | 下游五任务均值 | 0.5832 / 0.5737 / 0.5814 |
基线含 FA4(BF16/FP8)、cuDNN 三种精度、FlashInfer、Low Precision FA4、HW-FP4 FA4、PyTorch SDPA;各基线支持范围不同,提速按每负载最快可用者计。
泄露实验最有信息量。enwik8 上 6 层 GPT,带泄露的 4-bit 配置用自家核评估 loss 0.811–0.827,比 BF16 的 0.852 还低;换精确注意力评估立刻涨到 0.888–0.940,而无泄露训练的模型是 0.862。只看训练核自评的 loss,泄露完全不可见。前反向失配同样致命:对角处理正向 256、反向 128 queries,单步 dQ cosine 仍有 0.945,梯度范数却从 0.5 冲到 2×10^5;dS 尺度过大,零取整条目从 0.4% 涨到 29%,与 BF16 的 loss 差距从 0.0008 扩大到 0.049。对已有实现的探测中,cuDNN MXFP8 与 Low Precision FA4 默认配置在 1e-6 量级,HW-FP4 FA4 的 causal 路径 1.7e-1,FlashInfer NVFP4 KV 2.6e-1;只有 FLAMIA 的 exact 选项量到 0。
在 Blackwell 上预训练的团队可以直接换核:代码开源、带 TorchTitan 集成,1.1–1.4× 的注意力加速不碰任何超参。packed 原布局处理免掉 repacking,1B 全模型训练在 64k 行长下吞吐 96.8k tokens/s,padding 只有 84.3k,128k 时 padding 直接超显存。
比加速更耐用的是测量结论:单步 cosine 和自评 loss 都不足以判断低比特注意力可训不可训,causal 依赖探测加真实训练才算数。论文顺手用这套探测把主流实现量了一遍,几个默认路径都有非零泄露。
默认配置不是零泄露:块内 32 token 依赖和反向 dQ 累加器的序列级依赖仍在,完全消除的 exact 选项只覆盖部分 batched causal 配置,causal dQ scale 本身带 6–9% 训练调用开销。
训练验证规模偏小:1B 参数、100B tokens、2k 上下文。出发点是长上下文注意力开销,验证却停在 2k,更长上下文没有数据。下游单项有波动,ARC Challenge 从 0.3925 掉到 0.3618(MXFP8)、0.3677(NVFP4),均值接近掩盖了这点。
提速对的是每负载最快基线,各基线支持范围参差,不算严格同口径对照;B300 非 causal 九个负载里输两个,1.06× 的赢面也薄。全 4-bit 路径更慢且 loss 更高,4-bit 峰值算力在注意力上还兑现不了,反向不受 tensor core 限制。