JSSA、共同部分空間とGather-freeカーネルで長文脈デコードを高速化

長文脈LLMの推論では、注意計算そのものだけでなく、各デコードステップでKV-cacheから大量のデータを読み出すことが大きな負担になります。文脈が長くなるほどfull-KV attentionのメモリトラフィックは増加するため、必要なトークンだけを選ぶスパース注意は有力な方向です。しかし、トークン単位の選択には二つの課題があります。低次元のスコアリングではクエリとキーの関係を失う可能性があり、選択後のトークンが分散していると、gatherなどの読み出し処理が高速化を打ち消しかねません。

「Joint Subspace Sparse Attention for Efficient Long-Context Decoding」で提案されたJSSAは、この二つの課題をアルゴリズムとカーネルの両面から扱います。

手法のポイント

• 共有クエリ・キー部分空間。 JSSAは、クエリとキーを共通の低次元空間に写し、トークンスコアの保存誤差を小さくするように部分空間を学習します。単純な圧縮で関連性の構造を失うのではなく、選択に必要なランキング情報を残すことが狙いです。
• スペクトル解を使った最適化。 厳密なスコア保存目的は非凸ですが、著者らは扱いやすい上界となる周辺目的を構成し、閉形式のスペクトル解を導きました。論文のテスト条件では、厳密な目的をさらに直接最適化しても測定可能な改善は得られなかったとされています。
• Gather-freeの実行系。 選ばれたトークンがKV-cache上で散在すると、通常のgather処理が追加コストになります。JSSAはトークン単位のindexed-attentionカーネルを設計し、標準的なGrouped-Query Attention（GQA）のKV-cache配置で、KV headごとに異なる選択結果を扱えるようにしました。

評価結果

評価対象は、標準GQA、スライディングウィンドウとグローバル注意を組み合わせた構成、Multi-head Latent Attention（MLA）モデルです。概要によると、JSSAはfull-KV attentionに近い精度を保ちつつ、他方式で見られた生成長の大幅な増加や、アーキテクチャ依存の失敗を避けました。

NVIDIA B200で文脈長128Kを用いた場合、full-KV attentionに対して最大6.43倍のデコードスループットを達成したと報告されています。また、追加学習なしで、モデル固有のDSAと同程度の推論精度を示し、700B規模の大規模モデルでもDSAの96〜99%のデコードスループットに達したとされています。ただし、これらは論文が示した評価条件における数値であり、すべてのモデルやハードウェアにそのまま適用できるとは限りません。

意義と留意点

JSSAの重要性は、トークン選択の品質とメモリアクセスの効率を別々に最適化していない点にあります。スコアリングだけを改善しても不規則な読み出しは残り、カーネルだけを高速化しても誤った選択による精度低下は解決できません。両者を組み合わせることで、長文脈推論で支配的になりやすいKV-cache転送に直接対応しています。

一方、実際の効果は、モデルの注意構造、選択するトークン数、カーネル実装、GPUに依存します。提示された概要には、選択率や遅延の詳細な内訳、幅広いアクセラレータでの比較は含まれていません。導入時には、自社のモデルと文脈分布で再検証する必要があります。

出典：Hugging Face Daily Papers