SinkRouter：大規模言語モデルおよびマルチモーダルモデルにおける効率的な長文脈デコーディングのためのシンク認識ルーティング

概要: LLMやLMMにおける長文脈デコーディングでは、デコードの各ステップごとにGPUメモリから大量のKVキャッシュデータを読み込む必要があるため、注意（attention）は次第にメモリバウンドになります。既存の加速戦略はしばしば、有用な情報を捨ててしまう可能性のあるヒューリスティックな枝刈りに依存することで、効率と精度のトレードオフを行っています。さらに深い次元では、得点の高いトークンをすべて無差別に保持したり、初期トークンを不可欠なアンカーとして扱ったり、あるいはヒューリスティックなヘッドルーティングに頼ったりする傾向があり、注意（attention）サンク現象のメカニズムに対する理解が不十分であることを反映しています。本論文では、注意サンク現象が、学習中に構築される安定で到達可能であり、かつ誤差を制御可能な固定点に対応することを示します。この洞察に基づき、SinkRouter（サンクルーター）を提案します。これは学習不要の選択的ルーティング手法であり、サンク信号を検出して、そうであればゼロに近い出力を生成してしまう計算をスキップします。このメカニズムを実運用の加速へと落とし込むために、ブロック単位の分岐とSplit-K並列性を備えた、ハードウェアを意識したTritonカーネルを開発します。LongBench、InfiniteBench、CVBench、MileBench、MMVPといった多様な長文脈ベンチマーク群に対して、Llama-3.1-8B、Llama-3.1-70B、Yi-9B-200K、LLaVA-1.5-7B、LLaVA-1.5-13Bのような、テキストのみおよびマルチモーダルのバックボーンの両方を用いて、徹底的な評価を行います。これらの設定すべてにおいて、SinkRouterはデコーディング効率を一貫して向上させつつ、競争力のある精度を維持し、512Kのコンテキストで2.03倍の速度向上を達成します。