TileMix、アテンション行列内でFP16／INT8を混在させ、A100の4K prefillスループットでFlashAttentionの2.22倍を達成

TileMixは、混合精度の粒度をテンソル全体やkernel呼び出し全体から、アテンションスコア行列内のハードウェアに整合したtile groupへと細分化する。有効な各tile groupでは、1つのルーティングbitによってFP16を使用するか、INT8 Tensor CoreでQKスコアを計算してINT32で累積するかを決定する。両経路はスケール復元後、同じFP16 online softmaxの最大値、正規化係数、出力accumulatorを共同で更新する。これにより、疎なアテンションのように一部のtoken間インタラクションを直接削除することなく、完全な密アテンション接続を維持でき、モデルの再学習も必要としない。

ルーティング判定によってFlashAttention型の内部ループが遅くなるのを防ぐため、システムは同じquery tile rowの決定を64-bit bitmaskに圧縮し、bit shiftとmaskを使って定数時間で照会する。より長いシーケンスでは、1つのbitで隣接する複数のkey tileを制御できるため、ルーティングmetadataはKV head数とquery tile row数に応じてのみ増加する。公開されているTriton実装はこのほか、grouped-query attention、paddingなしの可変長batch、INT8 KV cacheインターフェース、LLaMAおよびQwenモデルへの注入コードをサポートする。

著者らは、単一のNVIDIA A100 40GB、batch size 8の環境で、量子化、スケール復元、ルーティング、データ転送、スケジューリングをすべてend-to-endのprefill時間に含めて評価した。LLaMA 3.2 3B-Instructへの4K入力では、tile groupの75%をINT8へルーティングするSpTrans構成が31.80K token/sを達成した。同じラッパー環境におけるFlashAttentionは14.33K token/sで、TileMixはその2.22倍に相当する。一方、全INT8経路は29.80K token/sだった。LongEvalと、中国語・英語データを含むLV-Evalの結果は、FP16 tileの空間的な配置とINT8のカバレッジ率がともに品質へ影響することを示している。混合構成は、全INT8で失われた精度を回復できる場合が多いものの、カバレッジ率を高めても、すべてのデータセットでFP16と同等になるわけではない。

この研究の実用的な価値は、普遍的かつ無損失な高速化率を主張することではなく、精度とスループットの間を調整できるインターフェースを提供する点にある。今後は、適応的ルーティングが追加コストに見合うかに加え、FP8、Hopper／Blackwell GPU、decode段階、より大規模なモデルでの性能を検証する必要がある。現時点のエビデンスはA100上のFP16／INT8経路に限られる。また、ライブラリにはcommitが2件しかなく、正式なreleaseも存在しないため、独立した再現実験と統合テストが依然として必要である。