💬 NLP

NOSA: Native and Offloadable Sparse Attention

本論文は、メモリ効率と生成品質の間のトレードオフを解決するためにCPU-GPU間のデータ転送を制約する、KVキャッシュ・オフローディングに特化した学習可能なスパース・アテンション機構であるNOSAを導入し、既存のベースラインに対して大幅なデコーディング・スループットの向上を実現している。

原著者： Yuxiang Huang, Pengjie Wang, Jicheng Han, Weilin Zhao, Zhou Su, Ao Sun, Hongya Lyu, Hengyu Zhao, Yudong Wang, Chaojun Xiao, Xu Han, Zhiyuan Liu

公開日 2026-01-30

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Yuxiang Huang, Pengjie Wang, Jicheng Han, Weilin Zhao, Zhou Su, Ao Sun, Hongya Lyu, Hengyu Zhao, Yudong Wang, Chaojun Xiao, Xu Han, Zhiyuan Liu

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、非常に高速だがメモリが極めて少ない、小さな高性能タブレットで、10万ページもある膨大な百科事典を読もうとしているところだと想像してください。そのタブレット（あなたのコンピュータのGPU）は驚くほど速いのですが、保持できるメモリはごくわずかです。一度に開いておけるのは、ほんの数ページ分だけです。残りの本は、部屋の反対側にある巨大で低速な図書室（あなたのコンピュータのCPU）に置かれています。

新しい文章を理解しようとするたびに、タブレットは図書室まで走って行き、必要な特定のページを掴み、それを持ち帰らなければなりません。もし単語ごとに何度も往復していたら、あなたは読み書きすることよりも、走っていることにほとんどの時間を費やすことになります。これが、長いテキストを処理しようとする現在のAIモデルが直面している問題です。データの転送による「交通渋滞」に陥ってしまうのです。

旧来の解決策（そしてなぜ失敗したのか）

1. 「ランダムに掴む」手法（学習なしのオフローディング）：
以前のいくつかの手法は、「図書室から、重要そうなページをランダムに数ページ掴んでおこう」という方法を試みました。

問題点： AIは本全体を読んで学習していましたが、テスト時には突然、ランダムな断片だけを読むよう求められました。これは、教科書全体を読んで期末試験の勉強をした生徒に対して、ランダムな一文だけを見てテストを受けるように指示するようなものです。生徒は混乱し、ミスを犯し、特に長い物語においては、回答の質が悪化してしまいます。

2. 「賢く掴む」手法（学習可能なスパース・アテンション）：
他の手法は、AIに賢さを教えようとしました。「どのページが重要かを正確に学習せよ！」と。

問題点： AIは正しいページを選ぶことは学習できましたが、図書室への「移動の効率性」については学習できませんでした。AIは、あちこちに散らばったページを選んでしまうかもしれません。すると、AIは依然として何度も往復しなければならず、その移動速度（データ転送）がボトルネックとなり、全体の動作を遅らせてしまいます。

新しい解決策：NOSAとNOSI

著者らは、NOSA（Native and Offloadable Sparse Attention）と、その相棒となるシステムであるNOSIという新しいシステムを提案しています。

NOSAを、AIの学生に対する新しいルールだと考えてください。

「近所付き合い」のルール： ランダムにページを選んだり、あちこちに散らばったページを選んだりする代わりに、AIは本の中で「近くにある」ページを選ぶように訓練されます。
例え話： あなたがミステリー小説を読んでいると想像してください。もしあなたが50ページ目にいるなら、次に必要になるのはおそらく49ページ目や51ページ目です。すぐに10,000ページ目が必要になることはまずありません。NOSAは、AIに自分の「近所（ネイバーフッド）」に留まるよう教えます。
メリット： AIが近くにあるページを選ぶため、図書室の棚からバラバラにページを拾うのではなく、一塊の「チャンク（塊）」としてまとめて掴むことができます。これにより、図書室への往復がはるかに速く、かつ頻度の少ないものになります。

NOSIは、これらのチャンクを運ぶために著者らが作り上げた、超効率的な配送トラックです。

旧式のトラックは、これらの特定のチャンクを運ぶには遅くて非効率でした。
NOSIトラックは、これらの「近所のチャンク」を物理的に可能な限り速く運ぶために特注されており、AIが待ち時間ではなく、読書に集中できるように設計されています。

研究結果

研究者らはこれらを様々なサイズのAIモデル（小・中・大）でテストし、以下の結果を得ました。

より高い品質： AIが「近所の」ページを選ぶように訓練されたため、「ランダムに掴む」手法のような混乱が起きませんでした。AIは長い物語や複雑な推論タスクをより正確に理解できました。
圧倒的なスピード： 図書室への往復回数を減らし、その移動を効率化したことで、システムは驚異的な速さになりました。
- 標準的な手法よりも最大5倍速くなりました。
- 最も優れた従来の「賢く掴む」手法よりも、ほぼ2倍速くなりました。
妥協なし： 彼らは、理解力を損なうことなく、このスピードを実現しました。AIは賢いまま、かつ高速なのです。

まとめ

この論文は、散らばったページではなく「情報の近所（ネイバーフッド）」に焦点を当てることで、AIに長い本を読ませる方法を提示しています。これにより、AIは低速なメモリへ行く代わりに、高速なメモリ内に留まることができ、より少なく、より効率的な移動が可能になります。その結果、より高価なハードウェアを必要とすることなく、AIは以前よりもはるかに速く、正確に膨大な量のテキストを処理できるようになりました。

技術要約: NOSA: Native and Offloadable Sparse Attention (ネイティブかつオフロード可能な疎なアテンション)

問題提起

大規模言語モデル（LLM）のデコーディング・スループットは、主にGPUメモリの制限、特にKey-Value（KV）キャッシュのストレージ要件によってますます制約を受けている。KVキャッシュのオフローディング（冗長なコンテキストをCPUに保持し、サブセットをGPUにフェッチする手法）はメモリ圧迫を緩和するが、既存のアプローチは以下の決定的なトレードオフに直面している：

学習フリーのオフローディング (Training-Free Offloading): ShadowKVやInfLLMのような手法は、学習なしで疎なサブセットを選択する。これは学習と推論のミスマッチ (training-inference mismatch) を引き起こす。すなわち、学習時に見られなかった疎なパターンが選択されることで、特に長文生成シナリオにおいて選択エラーが蓄積し、性能の大幅な低下を招く。
学習可能な疎なアテンション (Trainable Sparse Attention): 疎なパターンを学習する手法（InfLLMv2など）は、このミスマッチを回避できるが、効率的なオフローディングとは互換性がない。制約がない場合、これらの手法はCPUから任意のKVブロックをフェッチする可能性があり、その結果、スループットの向上を打ち消してしまうほどの膨大なPCIe通信オーバーヘッドが発生する。

現在のソリューションは、高いデコーディング効率（オフローディングによる）と、高い生成品質（学習との一貫性による）を同時に達成できていない。

手法

著者らは、KVキャッシュのオフローディングのために明示的に設計された学習可能な疎なアテンションメカニズムであるNOSA (Native and Offloadable Sparse Attention) と、それに特化した推論システムであるNOSIを提案する。

1. アルゴリズム: NOSA

NOSAは、デコーディングステップ間で疎なアテンションパターンが安定し、CPU-GPU間のデータ転送を最小限に抑えるために、学習中に局所性制約 (locality constraint) を導入する。

二段階選択メカニズム (Dual-Selection Mechanism): NOSAはKV選択を以下の2つのコンポーネントに分解する：
- クエリ依存型選択 ( $\Gamma_q$ ): クエリを使用して、特定の、潜在的に離れた情報を取り出し、（リコールを維持する）。
- クエリ非依存型選択 ( $\Gamma_e$ ): DMAに着想を得たエビクション・ヘッドを用い、クエリに依存しない重要度スコアを割り当てる。このコンポーネントは厳格なエビクション・ポリシーを強制する：あるトークンがステップ $t$ で選択されなかった場合、将来のステップ $t' > t$ においても選択されることはない。
局所性の保証 (Locality Guarantee): クエリ非依存型選択に固定予算 $k_e$ を割り当てることで、NOSAはステップ $t$ において選択されたKVブロックの大部分がステップ $t-1$ でも選択されていたことを数学的に保証し、 $\gamma(t) \geq k_e/k$ という最小限の局所性を保証する。これにより、それらのブロックをGPU上にキャッシュし続けることが可能になる。
ブロック単位の実装 (Block-Wise Implementation): PCIe帯域幅の利用を最大化するため、NOSAは個々のトークンではなくブロック単位で動作する。スコア計算には、二段階のプーリング戦略（サブブロック内での平均プーリング、それに続くブロックレベルでの最大プーリング）を採用している。
ED-DMA (Exp-Delayed DMA): エビクション・ヘッドにおける数値精度の問題を処理するため、NOSAは指数演算をアテンション計算ステージまで遅延させ、指数計算前のスコアに基づいてブロックを選択する。

2. システム: NOSI

著者らは、バニラの実装（例：Hugging Face Transformers）が過剰なカーネル起動オーバーヘッドを導入し、真の通信ボトルネックを隠蔽していることを指摘している。NOSIは、これらのボトルネックを露呈させ最適化するために構築されたカスタム推論システムである：

カーネル融合 (Kernel Fusion): エビクション・ヘッドの計算、QKV分割、およびプーリング操作を単一のTritonカーネルに融合し、起動オーバーヘッドを削減する。
最適化されたメモリレイアウト: KVキャッシュを最小限のブロックサイズ単位で構成し、カスタムCUDAカーネルを使用して、連続したメモリ要件を回避しながら、ブロックのスワップのための集合差分を効率的に計算する。
高帯域幅通信: ユニファイド仮想アドレス (UVA) を使用してカスタムのホスト・ツー・デバイス転送カーネルを実装し、標準的なPyTorchの実装における断片化の問題を克服して、ピークに近いPCIe帯域幅を実現する。

主な貢献

実現可能性の分析: 本論文は、学習可能な疎なアテンション（InfLLMv2など）は本質的な局所性を備えているものの、明示的な制約なしにはPCIe帯域幅の制限を克服するには不十分であることを特定した。
NOSAアルゴリズム: クエリ非依存のエビクション・ポリシーを通じて局所性制約を課すことにより、CPU-GPU間のKV転送量を明示的に制限する、新しい学習可能な疎なアテンションメカニズム。
NOSIシステム: NOSAに特化した推論システムであり、ソフトウェアのオーバーヘッドを排除することで通信を真のボトルネックとして浮き彫りにし、高スループットのオフローディングを可能にする。
包括的な評価: 1B、3B、8Bモデルを用いた広範なベンチマークにより、NOSAが学習フリーのオフローディング・ベースラインおよび他の学習可能な疎なアテンション手法を凌駕することを実証した。

実験結果

著者らは、一般的なタスク、長文コンテキスト入力、および長文生成推論タスクにおいて、1B、3B、8Bモデルを用いてNOSAを評価した。

パフォーマンス: NOSAは、長文コンテキストのベンチマーク（LongBench, HELMET）において、学習フリーのオフローディング・ベースライン（ShadowKV, InfLLM, ArkVale）および学習可能なベースライン（InfLLMv2, DMA）を一貫して上回る。学習フリーの手法が学習と推論のミスマッチにより性能を低下させる場面でも、NOSAは高いリコールと推論能力を維持する。
スループット:
- NOSAは、FullAttnと比較して最大5.04倍高いデコーディング・スループットを達成する（オフローディングによって可能になった大きなバッチサイズによる）。
- InfLLMv2と比較して最大1.92倍高いスループットを達成する（局所性制約による通信の削減による）。
- ShadowKVと比較して最大1.83倍高いスループットを達成する（疎なパターンの処理効率とシステム最適化による）。
長文生成: 長文生成の推論タスクにおいて、NOSAは低いパープレキシティを維持するが、学習フリーのオフローディング手法はエラーの蓄積とパープレキシティの爆発を示す。
アブレーション: 研究により、ブロック単位の選択パターンがPCIe効率に極めて重要であること、およびED-DMAエビクション・ヘッドの実装が最良のパフォーマンスをもたらすことが確認された。

意義と主張

本論文は、NOSAが「学習フリー」または「制約のない学習可能」なアプローチから、アルゴリズムとシステムの共同設計 (algorithm-system co-design) への必要な転換を象徴していると主張している。学習中にKV転送量を明示的に制約することにより、NOSAはメモリ効率（オフローディング）の必要性と、生成品質（学習の一貫性）の必要性の間の溝を埋める。

著者らは、NOSAが長文コンテキスト能力を持つLLMを、生成品質を犠牲にすることなく、メモリ制約のあるハードウェア上で展開することを可能にすると述べている。彼らは、デコーディング効率を向上させることで計算量とエネルギーコストを削減できることを強調しており、これがLLMアプリケーションのより広範な展開を促進する可能性があるとしている。研究を促進するために、コードは公開されている。