QFlash：桥接视觉中的量化和内存效率 Transformer Attention

FlashAttention通过平铺提高了效率，但其在线softmax仍然依赖浮点运算来实现数值稳定性，使得完全量化变得困难。我们确定了纯整数 FlashAttention 的三个主要障碍：（1）分块累积期间的规模爆炸，（2）GPU 上基于移位的指数运算效率低下，以及（3）需要统一比例进行整数比较的量化粒度约束。为了应对这些挑战，我们提出了 \textit{QFlash}，这是一种端到端整数 FlashAttention 设计，完全在整数域中执行 softmax 并作为单个 Triton 内核运行。在 ViT、DeiT 和 Swin 模型的 7 个注意力工作负载上，QFlash 比 I-ViT 实现了高达 6.73$\times$ 的加速，在 Swin 上实现了高达 8.69$\times$ 的加速，同时与 FP16 FlashAttention 相比降低了 18.8\% 的能耗，且不牺牲 ViT/DeiT 上的 Top-1 精度，并在每张量量化下在 Swin 上保持竞争力。我们的代码可在 https://github.com/EfficientCompLab/qflash 上公开获取。