vLLM Kimi K3 サポート

vLLM Kimi K3 サポート

vLLMは、2.8兆パラメータを持つオープンウェイトのマルチモーダル Mixture-of-Experts (MoE) モデルである Kimi K3 に対する、効率的な day-0 サポートを発表しました。この統合により、Kimi K3 の独自のハイブリッドアーキテクチャの高パフォーマンスなサービングが可能になり、16枚の NVIDIA GB300 NVL72 GPU 上で DSpark speculative decoding を使用した場合、ユーザーあたり最大 370 tokens/s を達成します。

Kimi K3 アーキテクチャの概要

Kimi K3 は、トークンごとに 896 個の専門家(experts)のうち 16 個をアクティブ化する 2.8T パラメータモデルです。以下の主要な革新技術を通じて、最大 100万トークンのウィンドウをサポートし、ロングコンテキストの効率性を実現するように設計されています:

  • Kimi Delta Attention (KDA): 線形アテンション層(固定サイズのリカレント状態を維持)と、グローバルな想起を維持するための周期的なフルアテンション層を交互に配置したハイブリッドスタック。
  • Attention Residuals (AttnRes): 標準的な残差蓄積を depth-wise attention に置き換え、学習された pseudo-queries を使用して前のブロックからの残差状態に重みを付けるメカニズム。
  • Stable LatentMoE: アクティベーションをより狭い潜在次元に投影してルーティングされた専門家の計算を行い、Quantile Balancing を利用してルーターのスコアの分位数に基づいて専門家を割り当てる設計。
  • Native MXFP4: バンド幅の要件を削減するために、MoE パスにネイティブな 4-bit 重みを使用。

Kimi K3 のための vLLM サービング最適化

Kimi K3 が標準的な Transformer アーキテクチャから逸脱しているため、vLLM はいくつかの特化したサービングメカニズムを実装しました:

ハイブリッド KV-Cache 管理

vLLM は、フルアテンション層用の paged KV ブロックと、KDA 層用のコンパクトなリカレント状態ブロックの両方を処理するために、単一のハイブリッド KV-cache マネージャーを使用します。このハイブリッド設計全体で prefix caching を有効にするために、vLLM は物理的な KDA 状態ブロックをきめ細かな prefix matching から分離し、共有プロンプトが KDA 状態と paged KV の両方を再利用できるようにしています。

融合カーネルとメタデータ最適化

  • Attention Residuals: vLLM は、融合された Triton および CUDA カーネルを使用して、depth-wise attention の logits、softmax、および hidden-state の集約を単一の操作で計算し、メモリトラフィックを削減します。
  • KDA Decode: 特化した CUDA カーネルが、因果的畳み込み(causal convolution)、リカレント更新、および RMSNorm を単一の起動に融合し、中間テンソルと繰り返される状態トラフィックを回避します。
  • KDA Prefill: vLLM は FlashKDA とその後のコミュニティによる最適化(Flash-Flash-KDA)を統合し、prefill フェーズを加速させました。
  • Metadata Builder: 専用の Kimi K3 KDA メタデータビルダーが、eager な PyTorch 操作を融合された Triton カーネルに置き換え、バッチサイズ 1 においてメタデータ準備のレイテンシを 96% 削減しました(870°s から 34°s へ)。

低レイテンシ計算

vLLM は、低バッチサイズ設定に対して skinnyGEMM を採用し、共有メモリのデータステージングをバイパスしてエンドツーエンドのレイテンシを約 10% 削減します。さらに、LatentMoE tail-fusion 最適化により、線形投影フェーズにおける冗長な計算を削減し、7°8% のエンドツーエンドのスピードアップを実現します。

本番環境での機能とパフォーマンス

DSpark による Speculative Decoding

vLLM は、ブロック拡散 speculative decoding アルゴリズムである DSpark をサポートしています。vLLM と TorchSpec でトレーニングされた DSpark speculator を使用することで、vLLM はシングルユーザーのリクエストに対して 3.14倍のスピードアップを達成し、スループットを 118 tok/s から 370 tok/s に向上させます。受理率は、低エントロピータスク(コーディング)のステップあたり 4.73 tokens から、高エントロピータスク(クリエイティブライティング)のステップあたり 2.61 tokens の範囲に及びます。

Prefill/Decode 分離 (Disaggregation)

高スループット環境向けに、vLLM は prefill/decode (PD) 分離をサポートしています。NIXL コネクタは、個別の prefill レプリカと decode レプリカの間で、トークンレベルの MLA キャッシュとリクエストレベルの KDA 状態の両方の転送を管理します。

Agentic キャッシュ保持ポリシー

ロングコンテキストの agentic ワークロードにおける KDA 状態のメモリフットプリントを管理するために、vLLM は 2 つの保持ポリシーを実装しています:

  • 間隔ベースの保持 (Interval-based retention): KDA 状態を固定間隔(例:32K トークンごと)およびプロンプトの境界で自動的にチェックポイント化します。
  • Marconi スタイルの選択的保持 (Marconi-style selective retention): デマンド駆動型のアプローチで、prefix の 2 回目のヒット時にのみ KDA 状態をキャッシュし、頻繁に共有される prefix のみがキャッシュ容量を消費するようにします。

ベンチマークとハードウェア要件

パフォーマンス指標

GB300 NVL72 GPU 上で、vLLM は以下のシングルユーザー decode スループット(バッチサイズ 1)を達成します:

設定 標準 Decode DSpark Speculative Decode
TP8 111 tok/s 331 tok/s
TP16 118 tok/s 370 tok/s

精度

vLLM 上の Kimi K3 は、最大推論設定において GSM8K で 0.976、GPQA-Diamond で 0.939、OCRBench で 0.889、MMMU Pro Vision で 0.818 のスコアを獲得しています。

ハードウェア要件

Kimi K3 のサービングには、最低 8× B300 (または GB300 NVL72) GPU、または 16× B200 GPU が必要です。NVIDIA (Hopper および Blackwell) と AMD (MI355X) の両方のハードウェアがサポートされています。

デプロイガイドの要約

Kimi K3 を実行するには、vLLM は以下の構成を推奨しています:

  • Prefix Caching: --enable-prefix-caching を介して明示的に有効にする必要があります。
  • MoE Backend: 分離型/エキスパート並列 (DEP) 環境には deep_gemm_mega_moe を、TP > 1 の場合は flashinfer_trtllm を使用してください。
  • All-to-all Backend: NVLink の場合は flashinfer_nvlink_one_sided を、RDMA の場合は deepep_v2 を使用してください。
  • Rust Frontend: VLLM_USE_RUST_FRONTEND=1 で有効にします。

Sources