FLAMIA: FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell

FLAMIA 是一族 Blackwell attention kernel，提供 MXFP8 与混合 NVFP4/FP8 训练路径，使用 token-block 缩放、因果泄漏防护、前向-反向数值一致性，以及 packed 文档的 GPU 侧调度；在合并的前向与反向 attention 调用上，于 B200 取得 1.10–1.40×、于 B300 取得 1.06–1.42× 加速。

长上下文下 attention 主导 Transformer 训练开销；尽管 Blackwell tensor core 以 BF16 的 2 倍和 4 倍速率运行 FP8 与 FP4 matmul，利用它们仍需解决量化开销、共享 scale 导致的因果信息泄漏，以及前向与反向之间的数值不匹配。硬件块缩放沿收缩轴进行，而 attention 的六个逻辑操作数参与具有不同收缩轴的乘积，因此朴素缩放会复制量化表示，并使 memory-bound 量化占据训练时间的很大一部分。共享 scale 还可能让未来 token 在存在因果 mask 的情况下改变更早输出，反向 pass 必须一致地重建前向的量化中间结果，否则训练会不稳定。Packed 变长文档还要求文档局部缩放与负载均衡，且无需重新打包。

FLAMIA 基于 warp-specialized FlashAttention-4 设计，并提供两条默认训练路径：MXFP8（E4M3 配 token-block scales）以及一条混合路径，其中 QKᵀ 及其重计算使用 NVFP4，每 16 个 head-dimension 元素一个 scale，其余乘积使用 E4M3。Token-block scaling 在 32 个连续 token 与完整 head dimension 上共享一个 2 的幂指数，使单个 8-bit payload 同时服务 K-major 与 MN-major 方向；dS 使用每 32×32 块一个指数。因果泄漏通过以下方式限制：操作数 scale 从不跨越超过 32 个位置，scale 指数向上取整以避免元素被 clamp，V 使用每 key scale 并在对角 tile 上做 K fold，另为 FP16 accumulator 提供可选 causal dQ scale。数值一致性通过以下方式强制：从相同的量化 Q 与 K 重计算 scores，复现 tiling、对角决策与指数求值，并用与 dP 相同的量化 dO 计算 D。Kernel 级工作将 Hadamard rotation、块统计与 code/scale 生成融合为一次 launch，在 log 域为 NVFP4 量化概率，将 QKᵀ 围绕 P̂V 拆分以保持 tensor core 忙碌，使用低精度 two-CTA 反向，并通过 GPU 构建的最长优先工作列表加 Cluster Launch Control 获取，就地处理 packed 文档。

• MXFP8 在两种 GPU 上的全部九个 causal packed-document 工作负载上，取得最低的实测 attention 训练调用延迟（含量化）。 §4.1 Performance on Packed Documents
• 4-bit NVFP4 kernel 在每个评估工作负载上都取得最低的 causal 反向延迟。 §4.1 Performance on Packed Documents
• 泄漏探针显示，per-tensor scales、沿 key 分组的 4-bit V、HW-FP4 FA4 causal 与 FlashInfer NVFP4 KV 会改变更早输出，而 FLAMIA 的 exact 选项使这些输出保持不变；cuDNN MXFP8 与 Low Precision FA4 会轻微改变它们。 §4.5 Causality and Training Stability9.8e-2 （原文中未核实）7.9e-2 （原文中未核实）2.8e-1 （原文中未核实）6.9e-6 （原文中未核实）1.7e-1 （原文中未核实）2.6e-1 （原文中未核实）1.1e-1 （原文中未核实）8.3e-2 （原文中未核实）
• 与 FP32 参考的余弦相似度显示，评估的 8-bit 配置接近 0.998–0.999，混合 NVFP4/FP8 在 causal O/dQ 上为 0.9920/0.9906，全 4-bit 为 0.9810/0.9456。 §4.2 Numerical Accuracy
• 在 1B 模型、100B token 预训练中，MXFP8 与混合 NVFP4/FP8 仅带来训练损失的边际增加，下游分数接近 BF16 基线。 §4.6 Training validation
• 推理期 bit-level 指数：Schraudolph 近似在余弦变化可忽略的情况下带来 MXFP8 加速，Mitchell-P 比 Direct-P 更频繁地复现 exact-rounding codes。 §4.4 Bit-Level Exponentials at Inference

• 默认缩放仍保留残余依赖：块内最多 31 个位置的 residue，以及序列范围的 dQ accumulator 依赖；移除这些依赖的 exact 选项目前覆盖受支持的 batched causal 配置，而 packed 文档、windows 与 all-4-bit 反向使用默认 scale（§5. Discussion and Conclusion; §4.5 Causality and Training Stability）。
• 移除 accumulator 依赖的 causal dQ scale 会在 32768 token 时使反向慢 8–11%，因此作为选项提供（§4.5 Causality and Training Stability）。
• All-4-bit 训练路径不是默认：它比将大多数操作数保留在 8 bit 的路径更慢，并达到略高的损失，其中量化六份副本占其反向时间的 45–52%（§3.4 MXFP8 and Mixed NVFP4/FP8 Paths; §4.3 Ablation Study）。
• Bit-level 指数不用于训练，因为反向必须精确重建前向的 P̂；在 causal 工作负载上，Mitchell-P 延迟约为 NVFP4 的 1.25–1.77×（§3.7 Bit-Level Exponentials for Inference; §4.4 Bit-Level Exponentials at Inference）。

• Token-block scaling 在 32 个连续 token 与完整 head dimension 上共享一个 2 的幂指数；由于该指数在 32×d 块内相同，codes 与 scales 在转置下仍有效，且单个 E4M3 payload 同时服务两个方向。 §3.1 Token-Block Scaling
• 仅靠 causal masking 不能阻止通过数据相关 scale 的泄漏；设计将任何操作数 scale 限制在至多 32 个位置，Q 与 K 使用 per-token NVFP4 scales，并将 scale 指数向上取整，使缩放后的最大值永不被 clamp。 §3.2 Causally Safe Scales
• 为前向-反向一致性，反向从相同的量化 Q 与 K 重计算 scores，复现 tiling、对角决策、指数求值与行最大值缩放，并用与 dP 相同的量化 dO 计算 D。 §3.3 Forward–Backward Numerical Consistency
• 在消融中，dS 条目被舍入为零的比例比报告范数误差更贴近训练退化，促使 scale 由每个 32×32 块的最大绝对值导出。 §3.3 Forward–Backward Numerical Consistency
• 将 Hadamard rotation、块统计与 code/scale 生成融合为一次 launch，在 4×4096 token 下把量化预处理从 MXFP8 的 789 降至 261 µs、从 NVFP4 的 0.64 降至 0.11 ms，并将 kernel launch 次数从 17 降至 1、从 11 降至 1。 §4.3 Ablation Study
• 在 log 域为 NVFP4 量化概率操作数，消除了逐条目除法与一次 scaling pass，报告为 NVFP4 前向最大的单项改进，达 14–19%。 §3.5.2 Softmax Throughput and MMA Scheduling
• Packed 文档使用 per-token scale metadata 就地量化，load warp 在 shared memory 中组装每个 tile 的 128×4 scale-factor atom，调度在每个 microbatch 于 GPU 上构建一次，并对偏斜或 causal 调用使用 CLC。 §3.6.2 Load-Balanced Scheduling
• 两种 bit-level 指数均不用于训练，因为反向必须精确重建前向的 P̂；在 4×4096 token 下，一致的 Mitchell-P 反向耗时 1.60–1.64 ms，而默认 NVFP4 路径为 1.00 ms。 §3.7 Bit-Level Exponentials for Inference