HiSparse：階層化メモリシステムによるスパースアテンションの加速

自己注意機構は、計算とメモリの二次コストのため、大規模言語モデル（LLMs）を長いコンテキストに拡張する際の主要なボトルネックとなっています。そのため、効率的な注意機構に対する関心が高まっています。その中でもスパースアテンションは特に注目されています。選択されたKVキャッシュのサブセットにのみ焦点を当てることで、スパースアテンションは強力なモデリング能力を保持しながら、コンテキストが成長するにつれて急激に増加する計算とI/Oコストを回避します。

しかし、スパースアテンション（通常はtop-k選択）はメモリ容量のボトルネックを解消できません。実際のアプリケーションでは、コンテキスト全体のKVキャッシュをGPUのHBMに保持し、迅速なアクセスを可能にする必要があります。これは、任意のデコードステップでわずかな部分のみがアクティブであっても同様です。このため、スパースアテンションは計算の制限ではなく容量の制限を受けやすく、実現可能なバッチサイズと全体的なスループットを制限します。これに対して、HiSparseは並列性を増加させた際にほぼ線形のスループット拡張を達成し、256の並列リクエストで基線スループットの3倍以上に達しました。低並列の場合、HiSparseは追加のI/Oがメモリの節約を上回るため、適度なオーバーヘッドを引き起こします。しかし、並列性の増加とメモリ圧力の支配に伴い、この利得は顕著になります。

HiSparseの設計

以前の作業HiCacheに基づき、私たちはHiSparseを提案します。これはこの制約を克服するための階層的なメモリシステムです。HiSparseは非アクティブなKVキャッシュエントリを積極的にホストメモリにオフロードし、GPUメモリの圧力を大幅に軽減します。一方、GPUのHBM上には頻繁にアクセスされるKV領域のホットデバイスバッファを維持し、クリティカルパス上のデータ移動を最小限に抑えます。これにより、より大きなデコードバッチが可能になり、スループットが向上し、より長いコンテキストに対応できます。下図はHiSparseのワークフローを示しています。プリフィルとデコードが分離された設定で描かれていますが、この設計は共置インスタンスにも適用されます。

効率的なSwap-inカーネル

このシステムのコアは専用のCUDAカーネルであり、効率的に：
(1)デバイスバッファ内のtop-kキャッシュ未ヒットを識別し、(2)LRU戦略により追放候補を選択し、(3)ページテーブルを更新して必要なエントリをホストからデバイスメモリに取得します。
下図はホットバッファサイズと追放戦略が未ヒット率に与える影響を示しています。より大きなホットデバイスバッファ（4096対2048スロット）とLRU追放戦略を使用することで、未ヒット数が大幅に減少し、クリティカルパス上でのswap-in遅延が直接的に低下します。

ベンチマークテスト

以下では、最先端のオープンモデルGLM-5.1-FP8に対する様々なシーケンス構成の結果を示し、長いコンテキストシナリオで最大5倍のスループット向上を実現しました。詳細はこちらをご覧ください。

# PD-disaggregation deployment (recommended) on two H20 nodes
# prefill instance:
python3 -m sglang.launch_server \
      --model-path "zai-org/GLM-5.1-FP8" --trust-remote-code --watchdog-timeout 100000 \
      --chunked-prefill-size 65536 --max-running-requests 480 --mem-fraction-static 0.8 \
      --tp-size 8 --dp-size 8 --enable-dp-attention --schedule-conservativeness 0.5 \
      --disaggregation-mode prefill \
      --disaggregation-ib-device mlx5_0,mlx5_1,mlx5_2,mlx5_3 \
      --dist-init-addr 127.0.0.1:5757 --nnodes 1 --node-rank 0

# decode instance:
python3 -m sglang.launch_server \
      --model-path "zai-org/GLM-5.1-FP8" --trust-remote-code --watchdog-timeout 100000 \
      --chunked-prefill-size 65536 --max-running-requests 480 --mem-fraction-static 0.85 \
      --tp-size 8 --dp-size 8 --enable-dp-attention \
      --load-balance-method round_robin --prefill-round-robin-balance \
      --kv-cache-dtype bfloat16 --nsa-decode-backend flashmla_sparse  \
      --disaggregation-mode decode --dist-init-addr 127.0.0.1:5757 \
      --disaggregation-ib-device mlx5_0,mlx5_1,mlx5_2,mlx5_3 --nnodes 1 --node-rank 0 \
      --enable-hisparse \
      --hisparse-config '{"top_k": 2048, "device_buffer_size": 6144, "host_to_device_ratio": 10}'


# PD-colocation deployment on a single 8xH200 instance
python3 -m sglang.launch_server \
      --model-path "zai-org/GLM-5.1-FP8" --trust-remote-code --watchdog-timeout 100000 \
      --chunked-prefill-size 65536 --max-running-requests 480 --mem-fraction-static 0.85 \
      --tp-size 8 --dp-size 8 --enable-dp-attention --disable-radix-cache \
      --enable-hisparse \
      --hisparse-config '{"top_k": 2048, "device_buffer_size": 4096, "host_to_device_ratio": 8}' 

今後の展望

HiSparseは現在DeepSeek Sparse Attention (DSA)を使用するモデルファミリーをサポートしており、DeepSeek-V3.2とGLM-5.1を含みます。実験的な機能として、私たちはその性能とモデルカバレッジの改善を続けることを期待しています。HiSparseは高並列シナリオでスループットを最大化するよう設計されていますが、top-kキャッシュ未ヒットによって引き起こされる追加のI/Oによっていくつかのオーバーヘッドも発生します。我々は、より良いオーバーラップによってこのオーバーヘッドを削減することを期待しており、Grace Blackwell（GB）システムのような新しいプラットフォームでのCPU–GPU帯域幅の向上がこの問題をさらに緩和することを信じています。

将来を見据えて、私たちの初期のHiCacheの作業の方向に沿って、この階層的メモリ管理方法を拡張し、ハイブリッドモデルを含むより広範な新興アーキテクチャをサポートする予定です。

謝辞

私たちはアリクラウドTairKVCacheチームとアントグループSCT推論チームの貴重な貢献に感謝しています。また、アリクラウドの蔡尚明、馬騰、凌星宇、およびSGLangコミュニティの許子藝からの寛大なサポートに感謝します。さらに、スタンフォードのChristos KozyrakisとKristopher Geda、そしてバイドゥ百歌AIチームからの深いフィードバックに感謝します。