効率的な低ビットMXFP推論のための対角タイル型ミックス精度注意

Abstract

トランスフォーマーベースの大規模言語モデル（LLM）は、幅広い実世界のタスクにおいて目覚ましい性能を示してきましたが、注意機構の二次的な計算量と、高精度演算におけるメモリ帯域制限のために、推論コストはいまだに過度に高いままです。本研究では、次世代GPUアーキテクチャ上の計算能力を活用し、microscaling floating-point（MXFP）データ形式を用いた、低ビットの混合精度アテンションカーネルを提案します。Diagonal-Tiled Mixed-Precision Attention（DMA）は、タイル化レベルで2種類の低ビット計算を取り入れており、モデル性能を損なうことなく高速かつ効率的な推論を実現するために、Tritonを用いてハードウェアレベルの並列性とメモリ効率を活用した、きわめて綿密に融合（fused）されたカーネルです。NVIDIA B200 GPUでの広範な実験評価により、提案カーネルは生成品質をほとんど劣化させることなく維持しつつ、一方でカーネル融合によって大幅な速度向上を達成することが示されています。コードをhttps://github.com/yifu-ding/MP-Sparse-Attn で公開します。