QFlash:量子化とメモリ効率を両立するビジョントランスフォーマーの注意機構

arXiv cs.LG / 2026/4/29

📰 ニュースDeveloper Stack & InfrastructureTools & Practical UsageModels & Research

要点

  • FlashAttentionはタイル化によって効率を高める一方、オンラインsoftmaxが数値安定性のために浮動小数点演算に依存しており、完全な整数量子化(完全量子化)を難しくしている。
  • 本論文では、整数のみのFlashAttention構築を阻む主要な障害として、タイル間の累積でのスケール爆発(オーバーフロー等)、GPU上でのシフトベースの指数演算の非効率、そして整数比較のために一様なスケールが必要になる量子化粒度の制約を挙げている。
  • QFlashは、softmaxを整数ドメインだけで完結させ、単一のTritonカーネルとして動作するエンドツーエンドの整数FlashAttention設計として提案されている。
  • ViT/DeiT/Swinの7つの注意ワークロードでの評価では、I-ViTに対して最大6.73×の高速化、Swinで最大8.69×の高速化を達成し、FP16 FlashAttentionと比べてエネルギー使用量を18.8%削減しつつ、ViT/DeiTではTop-1精度の低下はなく、Swinでもper-tensor量子化下で競争力を維持している。

要旨: FlashAttentionはタイル化によって効率を向上させますが、そのオンラインsoftmaxは数値的安定性のために依然として浮動小数点演算に依存しており、完全な量子化を困難にしています。整数のみのFlashAttentionに対する3つの主な障害を特定します: (1) タイル単位の加算中におけるスケールの爆発、(2) GPU上でのシフトベースの指数演算の非効率、(3) 整数同士の比較に一様なスケールを必要とする量子化粒度の制約です。これらの課題に対処するため、
\textit{QFlash}、すなわちソフトマックスを完全に整数ドメインで実行し、1つのTritonカーネルとして動作するエンドツーエンドの整数FlashAttention設計を提案します。ViT、DeiT、Swinモデルに由来する7つの注意(attention)ワークロードにおいて、QFlashはI-ViTに対して最大6.73\timesの速度向上を達成し、Swinでは最大8.69\timesの速度向上を達成しつつ、FP16 FlashAttentionと比べてエネルギー消費を18.8
%削減します。さらに、ViT/DeiTではTop-1精度を損なわず、Swinでも、per-tensor量子化の下で競争力のある性能を維持しています。私たちのコードは公開されています: https://github.com/EfficientCompLab/qflash。