让视觉变压器注意力计算全程用整数运算，速度提升近9倍且更省电。
QFlash: Bridging Quantization and Memory Efficiency in Vision Transformer Attention
- 设计整数域的FlashAttention，全程不用浮点数算softmax。
- 在7个ViT类模型上提速最高达8.69倍，能耗降18.8%。
- 适合追求高效推理的视觉模型部署场景，代码开源可复用。
FlashAttention通过分块加速计算，但其在线softmax仍依赖浮点数以保证数值稳定，难以实现全整数量化。我们发现三大障碍：(1) 分块累加时缩放值爆炸，(2) GPU上基于移位的指数运算效率低，(3) 量化粒度限制要求整数比较使用统一缩放系数。为此，我们提出QFlash，一种端到端整数化FlashAttention设计，可在单一Triton内核中完成全部softmax计算。在来自ViT、DeiT和Swin模型的七个注意力任务上，QFlash相比I-ViT最高提速6.73倍，Swin上最高达8.69倍，同时相较FP16 FlashAttention降低18.8%能耗，在ViT/DeiT上保持原始Top-1精度，Swin在每张量量化下仍具竞争力。代码已公开于https://github.com/EfficientCompLab/qflash。
原文摘要 · Abstract (English)
FlashAttention improves efficiency through tiling, but its online softmax still relies on floating-point arithmetic for numerical stability, making full quantization difficult. We identify three main obstacles to integer-only FlashAttention: (1) scale explosion during tile-wise accumulation, (2) inefficient shift-based exponential operations on GPUs, and (3) quantization granularity constraints requiring uniform scales for integer comparison. To address these challenges, we propose \textit{QFlash}, an end-to-end integer FlashAttention design that performs softmax entirely in the integer domain and runs as a single Triton kernel. On seven attention workloads from ViT, DeiT, and Swin models, QFlash achieves up to 6.73$\times$ speedup over I-ViT and up to 8.69$\times$ speedup on Swin, while reducing energy consumption by 18.8\% compared to FP16 FlashAttention, without sacrificing Top-1 accuracy on ViT/DeiT and remaining competitive on Swin under per-tensor quantization. Our code is publicly available at https://github.com/EfficientCompLab/qflash.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。