QFlash：视觉 Transformer 注意力中量化与内存效率的桥接
机器学习
2026-04-29 v1 人工智能
摘要
FlashAttention 通过分块实现效率提升，但其在线 softmax 仍依赖浮点运算以获得数值稳定性，使得完全量化困难。我们确定阻碍整数-only FlashAttention 的三个主要障碍：（1）块状累加期间的尺度爆炸，（2）GPU 上低效的基于平移的指数运算，（3）需要统一尺度的量化粒度约束进行整数比较。为解决这些挑战，我们提出 QFlash，一种端到端整数 FlashAttention 设计，完整在整数域中执行 softmax 操作，并作为单个 Triton kernel 实现。我们在来自 ViT、DeiT 和 Swin 模型的七个注意力工作负载上进行测试，QFlash 对 I-ViT achieves 最高 6.73 的加速，对 Swin 最高达到 8.69，同时在不牺牲 ViT/DeiT Top-1 准确率的情况下，将能源消耗比 FP16 FlashAttention 降低 18.8\%，在每个张量量化下对 Swin 仍保持竞争力。我们的代码已公开于 https://github.com/EfficientCompLab/qflash。
引用
@article{arxiv.2604.25306,
title = {QFlash: Bridging Quantization and Memory Efficiency in Vision Transformer Attention},
author = {Sehyeon Oh and Yongin Kwon and Jemin Lee},
journal= {arXiv preprint arXiv:2604.25306},
year = {2026}
}
备注
11 pages, 6 figures