Attn-QAT：量化感知训练恢复4位注意力质量

论文标题：Attn-QAT: Quantization-Aware Training Recovers FP4 Attention Quality
arXiv ID：2603.00040 | 发布日期：2026年3月 | 作者团队：Hao AI Lab @ UCSD

一句话总结

Attn-QAT 是首个系统研究 4 位（FP4）注意力量化感知训练的方法，通过识别并解决 FlashAttention 反向传播中的精度不匹配问题，实现了在视频扩散模型和大语言模型上恢复 BF16 级别的质量，同时在 RTX 5090 上比 SageAttention3 快 1.1-1.5 倍，在 B200 上快 1.39 倍。

核心问题

尽管 FP4 硬件（NVIDIA Blackwell 架构）已经支持，FP4 线性层已在生产中使用，但 FP4 注意力仍然导致显著的质量下降，主要原因包括：

原因	说明
FP4 动态范围极小	仅 19 个可表示值（E2M1 格式：-6 到 6），无法有效保留注意力动态
注意力激活分布重尾	与标准矩阵乘法相比，注意力有更多异常值，对数值精度更敏感

现有训练无关方法（如 SageAttention3 的 Q/K 平滑、双层量化）仍无法完全恢复质量。

FlashAttention 核心机制

1. 为什么不存储完整 P 矩阵？

朴素注意力实现需要 O(n²) 内存：

┌─────────────────────────────────────┐
│  标准实现:                           │
│  S = QK^T    → O(n²) 内存           │
│  P = softmax → O(n²) 内存           │
│  O = PV      → O(n²) 内存           │
│                                     │
│  总计: O(n²) —— 序列长度的平方！      │
│  n=8192 时，需要 ~500MB             │
└─────────────────────────────────────┘

FlashAttention 策略：实现 O(n) 内存复杂度

前向传播:     QK^T → S → softmax(P) → PV → O
                ↓
             保存 LSE (log-sum-exp 统计量)

反向传播:     从 LSE 重新计算 P → 计算梯度

2. 什么是 LSE？

LSE（Log-Sum-Exp）是 softmax 的数值稳定计算副产品：

原始公式:     P_i = exp(S_i) / Σ_j exp(S_j)

数值稳定版:
LSE_i = log(Σ_j exp(S_j))
P_i   = exp(S_i - LSE_i)

FlashAttention 实际计算过程：

m_i = max_j S_ij —— 行最大值
l_i = Σ_j exp(S_ij - m_i) —— 指数和
LSE_i = m_i + log(l_i) —— 行级 LSE

💡 关键洞察：只需保存 O(n) 的 LSE 向量，就能在反向恢复完整梯度计算。

3. 从 LSE 反向计算 P 的步骤

def recompute_attention_from_lse(Q, K, V, LSE, scale=1/sqrt(d)):
    # Q: [n_q, d], K: [n_k, d], V: [n_k, d], LSE: [n_q]
    
    # Step 1: 计算原始注意力分数
    S = torch.matmul(Q, K.transpose(-2, -1)) * scale  # [n_q, n_k]
    
    # Step 2: 减去 LSE 得到归一化后的 log 值
    # P_ij = exp(S_ij - LSE_i)
    log_P = S - LSE.unsqueeze(-1)  # broadcast LSE
    
    # Step 3: 恢复概率矩阵
    P = torch.exp(log_P)  # [n_q, n_k]
    
    return P

核心公式：P_ij = exp(S_ij - LSE_i)

验证：Σ_j P_ij = exp(-LSE_i) × Σ_j exp(S_ij) = exp(-LSE_i) × exp(LSE_i) = 1 ✓

量化带来的两个核心问题

问题一：前向/后向精度不匹配

标准 QAT（线性层）	FlashAttention + QAT
前向：fake quantize → BF16 GEMM → 输出	前向：FP4 fake quantize Q,K → 计算 P^F → 输出 O
后向：BF16 梯度 → 同样的 fake quantize → 更新权重	后向：❌ 直接用 BF16 重新计算 P（没有 fake quantize）
✅ 前后精度一致，梯度计算正确	→ 梯度基于高精度 P 计算，但前向用的是低精度 P^F → 梯度方向错误 → 梯度爆炸

# 前向 (FP4 模拟)
S_F = QF @ KF^T / sqrt(d)
P_F = softmax(S_F)  # 这是 fake quantize 后的结果

# 反向 (重新计算) — 错误！
S_BF16 = Q @ K^T / sqrt(d)  # 用 BF16 重新计算
P_BF16 = softmax(S_BF16)   # 这是真实 BF16 精度！

# P_F ≠ P_BF16 → 梯度方向错误！

问题二：FlashAttention 恒等式失效

FlashAttention 反向依赖关键恒等式来避免存储完整 P：

\[\mathbf{P}_i^\top \mathbf{dP}_i = \mathbf{dO}_i^\top \mathbf{O}_i\]

恒等式推导：

P_i^T dP_i = Σ_j P_ij · dP_ij
           = Σ_j P_ij · (dO_i^T V_j)        [因为 dP = dO · V^T]
           = dO_i^T Σ_j P_ij V_j
           = dO_i^T O_i

QAT 下的问题：

前向实际	反向假设
$\mathbf{O} = \sum_j \mathbf{P}^F \mathbf{V}^F$	$\mathbf{O} = \sum_j \mathbf{P} \mathbf{V}$

❌ 恒等式不再成立，梯度计算错误！

为什么线性层 QAT 没有这个问题？

线性层	FlashAttention Attention
不需要”反向重新计算”	需要从 LSE 重新计算 P
不依赖恒等式 `P^T dP = dO^T O`	依赖该恒等式节省内存
前后精度天然一致	前后精度容易不匹配

💡 核心原因：attention 是高度融合的算子，前向/后向紧密耦合，不能简单套用线性层的 QAT 方法。

Attn-QAT 的解决方案

解决方案一：反向重新计算 P 时强制 fake quantize

# 修复：反向重新计算时也要做 fake quantize！
# 反向传播
S = QF @ KF^T / sqrt(d)       # 用 fake quantize 后的 QF, KF
P_F = softmax(S)              # 模拟前向的量化效果
P_F_fake = fake_quantize(P_F) # 关键：强制 fake quantize！

解决方案二：前向额外计算高精度 O’ 用于反向

\[\mathbf{O}'_i = \sum_j \mathbf{P}_{ij} \mathbf{V}_j^F\]

# 前向传播
O_low = compute_attention(QF, KF, VF)  # 低精度输出，使用 P^F
O_high = compute_attention(QF, KF, VF) # 高精度输出 (用真实P而非P^F)
store(O_high)  # 仅用于反向梯度计算

# 反向传播
dS = P ⊙ (dO ⊙ O_high - P^T(dO ⊙ O_high))  # 使用O_high恢复恒等式

核心方法总结

问题	解决方案
P 前后精度不一致	反向重新计算 P 时强制 fake quantize
恒等式失效	前向额外计算高精度 O’ 用于反向

实验结果

视频扩散模型（Wan 2.1）

配置	VBench 总体质量
BF16	0.8335
FP4（无训练）	0.7968
SageAttention3	0.8203
Attn-QAT	0.8279 ✅

Attn-QAT 完全恢复 FP4 注意力带来的质量损失，在 99 个随机 VBench 提示的人工评估中与 BF16 无法区分。

大语言模型

继续预训练：Attn-QAT 恢复大部分质量，甚至在 WinoGrande 和 ARC-c 上超越 BF16
监督微调：Attn-QAT 可作为 BF16 的直接替代品

推理性能

平台	加速比
RTX 5090 vs SageAttention3	1.1x-1.5x
B200 vs FlashAttention-4	1.39x

加速来自消除 SageAttention3 额外的前处理（Q/K 平滑、双层量化）

结论

🔬 首次系统研究 4 位注意力 QAT，识别反向传播中精度不匹配的关键问题
🔧 两个核心修复：低精度重新计算 + 高精度辅助输出
⚠️ 训练方法与低比特内核需协同设计，不能简单套用线性层 QAT 方法
💡 FlashAttention 的”不存储 P、反向重新计算”设计是量化困难的根本原因