INT-FlashAttention: Enabling Flash Attention for INT8 Quantization

要約

大規模言語モデル (LLM) の基礎として、セルフ アテンション モジュールは、シーケンスの長さに関する二次時間とメモリの複雑さという課題に直面しています。
FlashAttention は、GPU メモリ階層を活用することでアテンションの計算を高速化し、メモリ使用量を削減します。
有望な研究の方向性は、FlashAttendant と量子化手法を統合することです。
このペーパーでは、Ampere GPU での FlashAttend の推論速度を大幅に向上させる、FlashAttend のフォワード ワークフローと互換性のある初の INT8 量子化アーキテクチャである INT-FlashAttend を紹介します。
完全な INT8 アクティベーションと一般行列乗算 (GEMM) カーネルを備えた INT-FlashAttendant プロトタイプを実装し、完全に INT8 入力を備えた最初のアテンション オペレーターとしました。
一般的なトークンレベルのトレーニング後の量子化フレームワークとして、INT-FlashAttendant は INT4 などの他のデータ形式とも互換性があります。実験結果では、INT-FlashAttendant は、標準の FlashAttendant と比較して推論速度が 72% 高速で、量子化誤差が 82% 小さいことが示されています。
FP16 および FP8 データ形式。

要約(オリジナル)

As the foundation of large language models (LLMs), self-attention module faces the challenge of quadratic time and memory complexity with respect to sequence length. FlashAttention accelerates attention computation and reduces its memory usage by leveraging the GPU memory hierarchy. A promising research direction is to integrate FlashAttention with quantization methods. This paper introduces INT-FlashAttention, the first INT8 quantization architecture compatible with the forward workflow of FlashAttention, which significantly improves the inference speed of FlashAttention on Ampere GPUs. We implement our INT-FlashAttention prototype with fully INT8 activations and general matrix-multiplication (GEMM) kernels, making it the first attention operator with fully INT8 input. As a general token-level post-training quantization framework, INT-FlashAttention is also compatible with other data formats like INT4, etc. Experimental results show INT-FlashAttention achieves 72% faster inference speed and 82% smaller quantization error compared to standard FlashAttention with FP16 and FP8 data format.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google