TileMix、ハードウェアタイル単位でLLM注意機構の精度を切り替え

導入

長い入力を処理するLLMのプリフィル段階では、生成を始める前に大量のトークンを一括処理しなければなりません。密な自己注意機構は、クエリとキーの組み合わせからスコアを計算するため、文脈長が伸びるほど計算量とメモリアクセスが大きくなります。INT8のような低精度化は効率改善に有効ですが、全ての計算を一律に低精度へ置き換えると、長文脈タスクの品質が低下する可能性があります。

TileMixは、トークン間の接続を削るのではなく、注意スコア行列のどの領域に高い数値精度が必要かを実行時に振り分けます。

タイルを単位にした混合精度

TileMixは注意行列を、アクセラレータの実行単位に合わせたスコアタイルへ分割します。さらに隣接するキー側タイルをグループ化し、各グループをFP16で計算するかINT8で計算するかを、コンパクトなビットマスクで表します。選択されたグループは、FlashAttentionに似た融合カーネルの中で処理されます。

この粒度は実装上重要です。個々の相互作用ごとに精度を指定すると、メタデータや分岐の負担が増えます。一方、あまり大きな単位で固定すると、精度を必要な場所だけに残す柔軟性が失われます。TileMixでは、一つのルーティングビットで複数の隣接キータイルを制御できます。これにより、ハードウェアに適したタイル形状を維持しながら、長い文脈でも管理情報を抑えられます。

FP16経路とINT8経路は、それぞれ別の注意計算を完結させるわけではありません。両方が共通のオンラインsoftmax状態を更新するため、異なる精度のタイルが一つの密な注意計算に参加できます。精度選択は固定的な量子化設定ではなく、カーネル内部で実行される空間的な判断になります。

疎な注意機構との違い

TileMixが振り分けるのは、相互作用を計算するかどうかではなく、その数値精度です。合法なタイルグループは処理されるため、密なトークン接続を保ちます。また、ルーティングのための再学習やモデル構造の変更も不要とされています。Grouped-query attention、可変長バッチ、INT8のキー・バリューキャッシュにも対応します。

実験から見える意味

論文はLongEval、LV-Eval、A100上のプリフィルベンチマークで、LLaMA、Qwen、Vicunaを評価しています。提供された概要によれば、一律INT8で失われた長文脈品質の一部を回復し、FP16と比べてプリフィルのスループットも改善しました。つまり、完全なFP16と一律低精度の二択ではなく、精度と効率の調整可能な境界を作る考え方です。

TileMixの意義は、混合精度をオフラインの量子化設定から、融合注意カーネル内の実行判断へ移した点にあります。ただし、密な注意機構の二次計算量そのものをなくすものではありません。実際の効果は、ルーティング方針、文脈長、GPU、モデル系列に左右されるため、公開実装と詳細な測定条件の確認が重要です。

出典：Hugging Face Daily Papers