Attn-QAT: 4位量化感知训练的注意力机制

Achieving reliable 4-bit attention is essential for end-to-end FP4 computation on emerging FP4-capable GPUs, yet attention remains a significant obstacle due to FP4's limited dynamic range and the heavy-tailed activations associated with attention. This research presents the first systematic study of 4-bit quantization-aware training (QAT) for attention mechanisms. It was found that a naive 'drop-in' QAT approach, which combines an FP4 forward pass with a high-precision Flash Attention (FA)-style backward pass, leads to training instability. Two key principles for stable FP4 attention were identified: (1) matching low-precision recomputation of attention scores in the backward pass, and (2) addressing implicit precision assumptions in FA's gradient calculation. Based on these insights, corresponding solutions are proposed.