StreamIndex：ストリーミングTop-kによるメモリ制限付き圧縮スパースアテンション
DeepSeek-V4のCompressed Sparse Attention (CSA)におけるメモリ問題を解決するStreamIndexが提案されました。
StreamIndexは、中間スコアテンソルを完全にマテリアライズせずに、チャンク化されたパーティション・マージTop-kドライバを使用することで、HBMの制限を超える大規模なシークエンス長を処理することを可能にしました。
実験結果から、StreamIndexは既存の手法と比較して、最大32倍のシークエンス長での実行を可能にし、セットオーバーラップリコールも高い精度を維持しました。
本手法は、インデクサのステップに特化しており、アテンションカーネルの高速化やエンドツーエンドのチェックポイント機能の提供は目的としていません。
大規模言語モデル（LLM）の処理能力向上を目指す研究が続いています。特に、長い文章を扱う際のメモリ消費量が課題となっていましたが、この度、DeepSeekが発表した「StreamIndex」は、この課題を根本的に解決する新しい技術です。これは、超長文処理におけるメモリ制約を克服し、LLMのスケールアップを大きく後押しするものと注目されています。
CSAとメモリの壁
DeepSeek-V3.2やV4で導入された「Compressed Sparse Attention (CSA)」は、LLMの処理速度を向上させるための技術です。これは、入力データ（キー）を圧縮し、クエリごとに最も関連性の高い上位k個（top-k）だけを抽出して注意機構（アテンション）を適用する仕組みです。しかし、従来のCSAの実装では、トップkを絞り込む前の段階で、巨大な中間スコアテンソル（[B, S, H_I, T] FP32）をメモリ上に一度に生成する必要がありました。例えば、シーケンス長Sが65,536の場合、この中間データは256GBにも達し、単一GPUのHBM（高帯域メモリ）の容量を遥かに超えてしまうという問題がありました。
StreamIndexによるメモリ制約の打破
このメモリ問題を解決するために開発されたのが「StreamIndex」です。これは、CSAパイプラインをTritonという高速計算フレームワークで実装したものです。StreamIndexの核となるのは、「チャンク分割・マージ型トップkドライバ」という仕組みです。このドライバは、巨大な中間テンソルをメモリ上に丸ごと生成することなく、データを分割しながら処理し、トップkの選定を段階的に行うことができます。実験では、従来の「メモリを確保してから処理する」方式ではS=65,536でメモリ不足（OOM）を起こしたのに対し、StreamIndexはS=1,048,576まで処理を延長し、ピークHBM使用量を6.21GBに抑えることに成功しました。これは処理範囲を32倍に拡張したことを意味します。
高い精度と実用的な性能
StreamIndexは、単にメモリを節約するだけでなく、高い精度を維持している点も特筆されます。従来の「メモリを確保してから処理する」方式の正解データ（ground truth）と比較した際のリコール（Recall）は、小さいシーケンス長では完全に一致しています。また、チャンクサイズやキータイルサイズ、トップkといった複数の設計パラメータを調整した実験でも、平均リコールは1.0000に達し、最小リコールも0.9980以上を維持しています。さらに、TileLangのアテンションカーネルと組み合わせた実証実験では、従来の方式がメモリ不足を起こしたのに対し、StreamIndexは18.56GBのピークHBMで1.97秒という実用的な速度で動作したとのことです。
まとめ
StreamIndexは、LLMの超長文処理における最大のボトルネックの一つであったメモリ制約を、革新的なデータ処理手法で克服しました。この技術は、より大規模で長いコンテキストを扱うLLMの開発において、非常に重要なブレイクスルーとなる可能性を秘めています。
原文の冒頭を表示（英語・3段落のみ）
View PDF
HTML (experimental)
Abstract:DeepSeek-V3.2 and V4 introduce Compressed Sparse Attention (CSA): a lightning indexer (a learned scoring projection over compressed keys) scores them, the top-k are selected per query, and a sparse attention kernel reads only those. Public CSA implementations materialize a [B, S, H_I, T] FP32 score tensor before the top-k reduction. With H_I=64 indexer heads and the V4-Flash compression ratio m=4, that intermediate is 256 GB at sequence length S=65,536, exceeding any single-GPU high-bandwidth-memory (HBM) budget. We present StreamIndex, a Triton implementation of the CSA pipeline whose central component is a chunked partition-merge top-k driver that never materializes the full intermediate. On synthetic-but-realistic V4-shaped inputs at the indexer-step (layer) level on a single NVIDIA H200, the materialize path runs out of memory (OOMs) at S=65,536 with V4-Flash dimensions; StreamIndex runs the same indexer to S=1,048,576 with 6.21 GB peak HBM, a 32x regime extension. Set-overlap recall against the materialize ground truth is bit-exact at small S where both fit; across three 5-point design-space sweeps (chunk size, key-tile size, top-k), mean recall rounds to 1.0000 with min recall at least 0.9980 in every cell. The chunked driver composes with TileLang's pipelined attention kernel: at S=262,144 with V4-Flash dimensions, the materialize indexer paired with TileLang attention OOMs while the chunked indexer paired with the same attention runs in 1.97 s at 18.56 GB peak. Our contribution targets the indexer step; we make no claim of a faster attention kernel or of real-checkpoint end-to-end behavior. Code: this https URL.
※ 著作権に配慮し、引用は冒頭3段落までです。続きは元記事をご覧ください。