技術要約：ハイブリッド長文脈モデルにおけるFull/GQAレイヤーのためのOracle誘導型疎なプリフィル（Oracle-Guided Sparse Prefill）

1. 問題提起

長文脈のプリフィル（prefill）は、大規模言語モデル（LLM）のサービングにおいて依然として重大な計算上のボトルネックとなっている。現代のアーキテクチャはハイブリッド設計（ローカル、疎、線形、および回帰的コンポーネントの組み合わせ）をますます採用しているが、Full AttentionまたはGrouped-Query Attention（GQA）レイヤーは、長距離の証拠を検索・融合するための主要なメカニズムとして残り続けていることが多い。GQAはKVキャッシュのメモリを削減するものの、シーケンス次元における高密度な探索そのものを排除するわけではない。

本技術レポートが取り組む核心的な診断的問いは、**「固定されたサポート粒度とtop-k予算の下で、既存のFull/GQAレイヤーにおける高密度なアテンション探索のうち、タスクレベルの挙動を維持するために実際に必要なものはどれくらいか？」**という点である。

現在の疎なアテンション手法（Quest、DuoAttention、MInference、NSAなど）は、解釈における課題に直面している。すなわち、疎化の導入に伴う品質の低下が、「予算不足」によるものなのか、「インデクサーの選択ミス」によるものなのか、「サポート共有の粗さ」によるものなのか、あるいは「ランタイムの実装上の差異」によるものなのかを区別することが困難である。制御されたリファレンス（参照基準）がなければ、疎なサポートの理論的な実現可能性と、インデクサーやランタイムシステムの実際の実装限界を区別することは難しい。

2. 手法

2.1 アテンション・マス Top-k オラクル（Attention-Mass Top-k Oracle）

著者らは、展開可能な加速手法としてではなく、制御された科学的計測器として**「アテンション・マス top-k オラクル」**を導入する。その目的は、「疎な予算の実現可能性」を「インデクサーのエラー」および「ランタイムの実装効果」から分離することにある。

• メカニズム: 各レイヤーおよびクエリ位置に対して、オラクルは完全な高密度アテンション分布を計算する。その後、クエリヘッド間でアテンションスコアを集計し、ヘッド平均化されたトークン重要度分布（$P_{\ell,t}$）を形成する。
• 選択: 指定されたtop-k予算（$k$）およびサポート粒度（クエリごと、または選択ブロック共有）に基づき、オラクルはこの分布から上位$k$個のトークンインデックスを選択する。
• 実行: モデルは、元のクエリ、キー、値を使用しつつ、選択されたサポートセットのみに限定してアテンション出力を再計算する。極めて重要な点は、オラクルはサポートの選択にのみ高密度アテンションを使用し、その後のモデルの軌跡（trajectory）は疎に計算されることである。
• 粒度: 本研究では2つの粒度を評価する：
  • クエリごと ($b_{sel}=1$): 各クエリ位置が独自のサポートを選択する。
  • 選択ブロック共有 ($b_{sel}=64$): システムコストを償却するために、64個の連続するクエリに対して単一のサポートセットを選択する。

2.2 KL蒸留によるヘッド集約型インデクサー（Head-Collapsed Indexer via KL Distillation）

オラクルに導かれ、著者らは凍結されたGQAバックボーンのための推論時インデクサーを導出する。

• 設計原理: バックボーンモデルは変更されない。補助的なインデクサーは、オラクルによって定義されたトークン重要度分布（$P_{\ell,t}$）を予測するように訓練される。
• 訓練目的: インデクサーは、高密度アテンション・マス分布からのKLダイバージェンス蒸留を通じて訓練される。
• アーキテクチャ: 著者らは**「ヘッド集約型（head-collapsed）」**設計を提案している。すべてのクエリヘッドまたはKVグループに対して個別のスコアリングチャネルを維持するのではなく、インデクサーはヘッド間で補助的なスコアリング空間を集約し、共有されたトークン重要度分布を予測する。これにより、投影アクティベーションが削減され、明示的な補助ヘッド次元の必要性が排除される。
• 推論: インデクサーは候補となるキーをスコアリングし、上位$k$個のインデックスを選択し、疎なGQAレイヤーは選択されたトークンに対してのみアテンションを再計算する。疎な予算（$k$）は、再学習なしで推論時に調整可能である。

2.3 証拠の分解

論文では、実現可能性とデプロイの準備ができていることを混同しないよう、結果を3つの明確な証拠タイプに分類している：

1. オラクル・ギャップ (Oracle Gap): 高密度アテンションとオラクルによる疎なアテンションとの差。これは、サポート・ファミリーが理論的に実行可能かどうかをテストする。
2. インデクサー・ギャップ (Indexer Gap): オラクルによるサポートを蒸留されたインデクサーに置き換えた際に生じる追加の劣化。
3. リアライゼーション・ギャップ (Realization Gap): ランタイムの制約（例：選択ブロック共有、カーネル融合）によって導入される追加の劣化。

3. 主な結果

3.1 オラクルの実現可能性（Oracle Gap）

• 高密度に近い性能: Qwenファミリーのチェックポイント（Qwen3-8B, Qwen3.5/3.6 27B/35B-A3B）全体において、クエリごとのオラクルは、評価された最長文脈（最大128K）において高密度アテンションの1パーセントポイント以内に収まっている。
• 予算への感度: Qwen3.5-9Bの4Kから100Kまでのスイープにおいて、構成されたtop-kスケジュール（256–2048）の下で、オラクル・ギャップは高密度アテンションから0.48ポイント以内に留まり、疎化率は87.9%から96.0%であった。
• 粒度の影響: クエリごとのオラクルは堅牢であるが、選択ブロック共有（$b_{sel}=64$）のオラクルは「崖（cliff）」現象を示す。128Kにおいて、高い予算（top-k=2048）では高密度に近い品質を維持できるが、低い予算（top-k=256）では、検索重視のタスクで大幅な崩壊が見られる。これは、サポートの共有には長さに応じた予算のスケーリングが必要であることを示唆している。

3.2 蒸留インデクサーの品質（Indexer Gap）

• 品質の維持: Qwen3.5-0.8BおよびQwen3.5-9B用に個別に蒸留されたインデクサーを使用した場合、著者らはRULER、VideoMME、MMLongBench-Doc、およびLongBench-v2の保持された検証スライス（16K/32K文脈）において、実質的な集約的劣化は認められないと報告している。
• マクロ・ギャップ: 報告されたマクロ・ギャップ（蒸留後マイナス高密度）は、Qwen3.5-0.8Bで**+2.04**、Qwen3.5-9Bで**+1.13**であった。著者らは、これらの正のギャップを、高密度アテンションに対する本質的な改善ではなく、評価解像度内での品質維持として扱っている。
• クエリごと vs 融合: 純粋なRULER-32K検証において、クエリごとの蒸留パスは高密度から0.6ポイント以内であったが、融合（選択ブロック共有）パスは5.0ポイントのギャップを示した。これは、システムフレンドリーな償却に伴うコストを浮き彫りにしている。

3.3 ランタイムの余地（TTFT）

• 蒸留インデクサーによる高速化: シングルカードでの測定により、蒸留インデクサーを用いた疎なプリフィルのTTFT高速化が示されている：
  • Qwen3.5-0.8B (NPU): 256Kコンテキストにおいて最大1.71倍の高速化。
  • Qwen3.5-9B (GPU): 256Kコンテキストにおいて最大1.93倍の高速化。
    高速化はコンテキスト長とともに増大する。これは、非アテンションレイヤーとオーバーヘッドの相対的なコストが減少するためである。
• 未学習インデクサーのストレス・テスト: ランダムに初期化された（未学習の）インデクサー重みを用いて疎なパスを実行すると、さらに高い潜在的な余地（128KにおいてNPUで最大3.44倍）が明らかになった。これは、疎なランタイムパス自体は効率的であることを示しているが、出力品質はこの構成では検証されていない。

4. 重要性と主張

本論文は、完全に展開されたロスレスな疎システムであるという主張ではなく、**「診断的技術レポート」**としての立場をとっている。その主な貢献は以下の通りである：

1. 問題の再定義: 疎なプリフィルをGQAレイヤーにおける**「オラクル誘導型の還元可能性の問題」**として捉え直している。高密度アテンション・マスを用いて理論的に可能なことを確立することで、実用的な疎なシステムにおけるエラーの源を特定している。
2. オラクルの定義: アテンション・マス top-k オラクルは、予算の実現可能性とインデクサーの性能およびランタイムの実装を分離するための「制御されたリファレンス」として機能する。
3. 経験的証拠: QwenファミリーのGQAチェックポイントにおいて、高密度アテンション・マスが、極端な長さ（128K–256K）においてもタスクレベルの挙動を維持するコンパクトでクエリ依存のサポートを特定できるという最初の証拠を提供している。
4. 手法の分離: 論文は、以下の3つの異なる主張を明示的に切り離している：
   • オラクルの実現可能性: 疎なサポートは実行可能である（オラクル行によって証明）。
   • 蒸留インデクラーの品質: 凍結されたバックボーンのインデクサーは、品質損失を最小限に抑えてこのサポートを近似できる（蒸留行によって証明）。
   • ランタイムの余地: 疎なサービングパスは、大幅なTTFT削減を提供する（ストレス・テスト行および蒸留TTFT行によって証明）。

制限事項と範囲:
著者らは、これが「最初のリリース」であることを述べ、その主張には謙虚である。彼らは明示的に以下のことを述べている：

• すべての指標が全く同じ構成で測定された、単一の「完全に一致した高密度–オラクル–蒸留品質・レイテンシの境界線」は存在しない。
• 蒸留されたインデクサーは主に16K/32Kで検証されており、最大のTTFT高速化は、ペアとなる品質チェックなしに128K/256Kで測定されている。
• ヘッド集約型インデクサーのデザインは、システム指向の選択として提示されており、ヘッド拡張型に対するアブレーションによる最適性の証明はなされていない。
• 現在の証拠はQwenファミリーに集中しており、他ファミリーへの検証は今後の課題とされている。

要約すると、本論文は、GQAレイヤーにおける長文脈の検索において**高密度アテンションはしばしば「過剰（overkill）」**であり、蒸留されたヘッド集約型インデクサーを用いることで、サポートの粒度と予算を慎重に管理すれば、品質を維持しつつ大幅なTTFTの高速化を実現できることを示している。