TileMix：タイル中心の混合精度アテンションによるLLM推論高速化

要旨

大規模言語モデル（LLM）における長コンテキストのプリフィルは、高密度セルフアテンションがクエリ・キー間のスコアを二次的に計算するため、多大な計算量とメモリトラフィックを引き起こす。既存手法は、一律の低精度パスを用いるか、トークン間の相互作用を選択するかのどちらかであり、融合高密度アテンション内でハードウェア整列スコアタイルにわたる空間的精度ルーティングを扱っていない。我々は、融合高密度アテンション内のスコアタイル群に対して数値精度を実行可能な空間的決定とする、タイル中心の精度ルーティングカーネルであるTileMixを提案する。TileMixは、アテンション行列をハードウェア整列スコアタイルに分割し、ルーティング決定をコンパクトなビットマスクにパックし、各タイルグループをFP16またはINT8スコア計算に振り分ける。その間、両方のパスは共有されたオンラインソフトマックス状態を更新する。スケーラブルな精度グループ化により、各ルーティングビットが複数の隣接キータイルを管理でき、長いコンテキストでもハードウェア整列の計算タイルとコンパクトなメタデータを維持する。すべての正当なタイルグループをルーティングすることにより、TileMixは高密度トークン接続性を維持し、学習を必要とせず、グループ化クエリアテンション、可変長バッチ、INT8キー/バリューキャッシュをサポートする。LLaMA、Qwen、VicunaにおけるLongEval、LV-Eval、A100プリフィルベンチマークにわたって、TileMixは一律INT8で失われた長コンテキスト品質を回復し、FP16と比較してプリフィルスループットを向上させ、モデルファミリー間で制御可能な精度効率フロンティアを実現する。実装はhttps://github.com/HanzhiZhang-Ulrica/TileMixで利用可能である。

English

Long-context prefill in large language models (LLMs) incurs substantial computation and memory traffic because dense self-attention computes quadratic query-key scores. Existing methods either use a uniform low-precision path or select token interactions, leaving spatial precision routing over hardware-aligned score tiles outside fused dense attention. We introduce TileMix, a tile-centric precision-routing kernel that makes numerical precision an executable spatial decision over score-tile groups within fused dense attention. TileMix partitions the attention matrix into hardware-aligned score tiles, packs routing decisions into compact bitmasks, and dispatches each tile group through FP16 or INT8 score computation while both paths update a shared online-softmax state. Scalable precision grouping lets each routing bit govern multiple adjacent key tiles, preserving hardware-aligned compute tiles and compact metadata at long contexts. By routing all legal tile groups, TileMix preserves dense token connectivity, requires no training, and supports grouped-query attention, variable-length batches, and INT8 key/value caches. Across LongEval, LV-Eval, and A100 prefill benchmarks on LLaMA, Qwen, and Vicuna, TileMix recovers long-context quality lost under uniform INT8 and improves prefill throughput over FP16, yielding a controllable accuracy-efficiency frontier across model families. The implementation is available at https://github.com/HanzhiZhang-Ulrica/TileMix.