vLLM DeepSeek V4 サポート: 効率的な長文コンテキスト注意機構

vLLM は DeepSeek V4 モデルファミリーのサポートを発表しました。DeepSeek-V4-Pro(1.6T パラメータ)と DeepSeek-V4-Flash(285B パラメータ)を含みます。これらのモデルは、メモリオーバーヘッドと計算コストを大幅に削減しながら、最大 100 万トークンにスケールするよう設計された特殊な長文コンテキスト注意機構を利用します。

DeepSeek V4 アテンションアーキテクチャ

DeepSeek V4 は、KV キャッシュメモリの線形増加と長文コンテキスト注意の高い計算コストに対処するため、4 つの主要なアーキテクチャ上の革新を導入しています。

  • Shared Key and Value Vectors: キーとバリューのベクトルを共有することで、モデルはメモリを 2 倍削減します。正確性を保つために、inverse RoPE 操作がアテンション出力に適用され、平行移動不変性が確保されます。
  • KV Cache Compression: モデルは複数トークンにわたって KV キャッシュを圧縮し、4 倍から 128 倍のメモリ削減を 2 つの方法で実現します:
    • c4a: KV キャッシュを約 1/4 に圧縮します。1 つの圧縮トークンは、ストライド 4 で 8 つの非圧縮トークンの加重和を表します。
    • c128a: KV キャッシュを約 1/128 に圧縮します。1 つの圧縮トークンは、ストライド 128 で 128 の非圧縮トークンの加重和を表します。
  • DeepSeek Sparse Attention (DSA): 圧縮後でも計算量を抑えるため(例: c4a を使用した 1M シーケンスで 250k トークン)、DSA は上位 k の圧縮トークンのみに注意を向けます。
  • Short Sliding Window: サイズ 128 のスライディングウィンドウを非圧縮トークンのローカル情報に使用し、クエリトークンが圧縮境界に到達する前にローカルコンテキストにアクセスできるようにします。

これらの最適化により、メモリ効率が大幅に向上します。DeepSeek V4 の 1M コンテキストシーケンスは、シーケンスあたりわずか 9.62 GiB の KV キャッシュ(bf16 使用)で済み、61 層の DeepSeek V3.2 スタイルスタックが必要とする 83.9 GiB の約 8.7 倍小さくなります。インデクサキャッシュに fp4、アテンションキャッシュに fp8 を使用すると、さらに約 2 倍サイズが削減されます。

vLLM で DeepSeek V4 を実装するには、異種アテンションタイプとメモリ管理に関する複雑なシステム課題を解決する必要がありました。

KV キャッシュメモリ管理

vLLM は KV キャッシュをコンパクトかつ効率的に保つために、3 つの戦略を採用しています:

  1. Single Logical Block Size: vLLM はすべての圧縮層に対して論理ブロックを 256 ネイティブトークン位置 に固定します。これにより、レイヤーが c4a(ブロックあたり 64 圧縮エントリ)であろうと c128a(ブロックあたり 2 圧縮エントリ)であろうと、アロケータはスロットマッピングとプレフィックスヒット検出に一貫した単位を使用できます。
  2. Compressor State as Sliding Window: ローリング残差(C4 は 8 トークン、C128 は 128 トークン)の複雑なサイドバッファを回避するため、vLLM はコンプレッサ状態をスライディングウィンドウ KV として扱います。これにより、システムはプレフィックスキャッシュと分散プリフィルの既存抽象を再利用できます。
  3. Unified Page Sizes: ブロックサイズと圧縮率を慎重に選択することで、vLLM は 5 段階のキャッシュスタックを 3 つのページサイズバケット に統合します。これにより、プール間の断片化が解消され、実行時の再分割が不要になります。

GPU 計算飽和

GPU の利用率を最大化し、HBM の往復回数を減らすため、vLLM は複数のカーネル融合とマルチストリーム分割を実装しました:

  • Kernel Fusions:
    • Compressor + RMSNorm + RoPE + cache insertion: 要素ごとのステージを 1 つのカーネルに融合し、1.4〜3 倍の速度向上を提供します。
    • Inverse RoPE + fp8 quant: これらの操作を融合して HBM の往復を回避し、2〜3 倍の速度向上を実現します。
    • Fused Q norm + KV RoPE + K insert: クエリと非圧縮 SWA キーの処理を水平に 1 カーネルに融合し、10〜20 倍の速度向上を提供します。
  • Multi-stream Partitioning: vLLM は CUDA ストリーム間で独立した操作を重ね合わせます。c4a レイヤーでは、インデクサパイプラインが別ストリームで実行され、メイン KV 圧縮および SWA トークン挿入と並列に動作し、低バッチサイズでエンドツーエンドのレイテンシを 5〜6% 削減します。

デプロイとハードウェアサポート

DeepSeek V4 は NVIDIA Hopper および Blackwell アーキテクチャでサポートされています。vLLM はシングルノード展開向けに最適化された Docker 設定を提供します:

  • DeepSeek-V4-Pro: 8xB200 または 8xB300 GPU 向けに最適化。
  • DeepSeek-V4-Flash: 4xB200 または 4xB300 GPU 向けに最適化。

両構成は fp8 KV キャッシュ、ブロックサイズ 256、deepseek_v4 トークナイザーと推論パーサーを使用します。DeepGEMM MegaMoE カーネルやページングプリフィルカーネルなどのさらなる最適化が現在開発中です。

Sources