vLLM 上の Reflection Beam 1.0: デイ 0 サービング、エージェント最適化カーネル、DFlash2 投機的デコーディング

vLLM は、Reflection の Beam に対してデイ 0 サービングサポートを追加しました。これは 501B パラメータの mixture-of-experts モデルで、トークンあたり 23B のアクティブパラメータを持ち、NVIDIA Blackwell ハードウェア向けに特別に最適化されています。この結果が重要なのは、本番環境向けの完全なサービング機能（推論、ツール呼び出し、構造化出力、テンソル/データ/エキスパート並列処理）と、コーディングエージェントが自律的に最適化したカーネルスタックを組み合わせ、バッチ 1 デコードを vLLM の最初の実装よりも 2.71 倍、プリフィルを 1.87 倍高速化したからです。

単一リクエストは、Beam の DFlash2 ドラフトヘッドを使用して TP8 上で毎秒 708 トークンに達し、受容長は 3.7 です。適応的検証により、高並列時にはスループットがさらに 10% 向上します。

このモデルは、Docker イメージ vllm/vllm-openai:reflection-beam で提供できます：

docker run --gpus all --ipc=host -p 8000:8000 \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  vllm/vllm-openai:reflection-beam \
  reflection/Beam-1.0-501B-A23B-NVFP4 \
  -tp 4 \
  --max-num-seqs 128 \
  --reasoning-parser beam \
  --tool-call-parser beam --enable-auto-tool-choice \
  --kv-cache-dtype fp8 \
  --attention-config '{"backend":"FLASHINFER"}' \
  --kernel-config '{"enable_flashinfer_autotune":false}' \
  --speculative-config '{
    "method": "dflash",
    "model": "reflection/Beam-1.0-501B-A23B-DFlash2",
    "num_speculative_tokens": 7,
    "attention_backend": "FLASH_ATTN"
  }'

Beam モデルアーキテクチャ: 52 層、ハイブリッドアテンション、DeepSeek-V3 スタイルの MoE

Beam は、隠れサイズ 5120 の 52 層 MoE デコーダーです。その MoE ルーティングは DeepSeek-V3 に従います。アーキテクチャの特徴的な選択は、アテンション層と残差ストリームの正規化方法にあります。

アテンション: スライディングウィンドウ層と NoPE フルアテンション層

Beam はアテンションを 2 つの方式に分割します。43 層は RoPE を使用して最後の 512 トークンに対するスライディングウィンドウアテンションを使用します。残りの 9 層（6 層ごと、および最後の層）は、位置エンコーディングなし（NoPE）でフルコンテキストにアテンドします。

RoPE がない場合、キー数が増えるにつれてアテンションが平坦化するため、Beam は NoPE 層に SSMax を適用します。SSMax は各クエリを (1 + p/2048)^0.1 でスケーリングします。ここで p は 1 から始まる位置で、コンテキストが長くなるにつれてソフトマックスが鋭くなります。

Beam はグループ化クエリアテンションを使用し、48 のクエリヘッド、8 の KV ヘッド、ヘッド次元 128 を持ちます。Q と K は RoPE または SSMax の前にヘッドごとに RMS 正規化されます。出力プロジェクション（o_proj）の前に、各ヘッドの出力は sigma(gate + bias) で乗算されます。ここでゲートは独自の全幅プロジェクションと学習されたヘッドごとのバイアスから得られ、ヘッドをトークンごとに弱めたりオフにしたりできます。

アテンション設計により KV キャッシュは小さく保たれます。フルアテンション層は 9 層のみで、シーケンス全体の KV を保持します：トークンあたり 9 * 2 * 8 * 128 個の値で、BF16 では 36 KiB、FP8 では 18 KiB です。完全な 1M トークンコンテキストは、シャーディング前で約 18 GiB の FP8 KV キャッシュを消費します。

MoE と残差ストリーム: FP32 ルーター、サンドイッチ正規化、深さスケーリング残差

Beam は 256 のルーティングエキスパートを持ち、d_ff は NVFP4 で 2,560 です。ルーティングは DeepSeek-V3 と一致します：シグモイドスコアと補正バイアスで 8 つのエキスパートグループのうち最良の 4 つを選び、次に上位 8 エキスパートを選び、再正規化された重みは 2.5 でスケーリングされます。DeepSeek-V3 とは異なり、Beam には共有エキスパートがありません。層 0 から 2 は MoE の代わりに d_ff 20,480 の高密度 SwiGLU MLP を使用します。

ルーターは完全に FP32 です：5120 x 256 のルーター重み、累積、ロジットはすべて FP32 で実行されます。

2 つの正規化の選択が残差ストリームを形成します。すべてのサブレイヤーにはその前後に RMSNorm があります（「サンドイッチ正規化」）。n 番目の残差加算は h + x ではなく h <- sqrt((n-1)/n) * h + x/sqrt(n) を計算するため、単位分散のサブレイヤー出力では残差ストリームは深さが増しても単位分散を維持し、成長しません。

エージェント駆動の最適化によりバッチ 1 デコードが 2.71 倍高速化

vLLM の最初の Beam 実装は最適化されておらず、ほぼ純粋な PyTorch でした。その後、コーディングエージェントが自律的に最適化しました：モデルをプロファイリングし、カーネルを書き、各変更をベンチマークしました。TP8 では、7 つの変更によりバッチ 1 デコードがステップあたり 8.95 ms から 3.30 ms に、8 x 2K プリフィルが 328 ms から 176 ms に短縮されました。

7 つの変更は次のとおりです：

  1. ノルム、残差、ノルムの融合（-40% デコード）。104 のサブレイヤー境界それぞれでのポストノルム、スケーリング残差加算、次のプレノルムが 1 つの Triton カーネルとして実行され、FP32 演算と参照実装の BF16 丸めポイントを維持します。
  2. Q/K ノルム、RoPE/SSMax、出力ゲートの融合。ヘッドごとの Q/K RMSNorm と RoPE（または SSMax）が 1 つのカーネルとして実行され、シグモイドゲートは別のカーネルとして実行されます。
  3. テンソル並列オールリデュースをノルムカーネルに融合。FlashInfer のマルチキャストオールリデュースに基づきます。MoE の後、同じカーネルが上位 8 エキスパートの出力も結合します。
  4. FP32 ルーター用の専用 GEMM カーネル。ルーターは以前は汎用パスにフォールバックしていました。vLLM の小バッチ FP32 ルーターカーネルは、Beam の 5120 x 256 形状をサポートし、FP32 重み、累積、ロジットを維持します。
  5. QKV とゲートプロジェクションの低レイテンシ GEMM への統合。Q、K、V、ゲートの重みはロード時に 1 つの行列にパックされ、4 つのプロジェクションが 1 つの GEMM として実行されます。小バッチでは低レイテンシの CuTe DSL GEMM または FlashInfer の TinyGEMM を使用します。
  6. 融合 KV キャッシュ更新。Q/K ノルムと RoPE カーネルが K と V を KV キャッシュに書き込み、別のキャッシュ書き込みカーネルを排除します。
  7. 3 つの高密度層のアクティベーションを 1 つのカーネルで実行。

レイテンシ測定は GB200 上の TP8 モデルフォワードレイテンシとして、バッチ 1 デコードと約 100K コンテキストでの 8 x 2,048 トークンプリフィルに対して行われました。

融合デコードグラフは 1 つの MoE 層を 13 カーネルで実行

これらの変更後、TP8 のバッチ 1 デコードでは 1 つの MoE 層が 13 カーネルとして実行されます。そのうち 6 つは Beam 用に作成、融合、または拡張されたもので、Beam 固有の小さな操作はすべて隣接する操作に統合され、両方のオールリデュースはノルムカーネル内にあります。

残りの 7 つのカーネルはライブラリから来ています：TRTLLM-gen アテンション、cuBLAS 出力プロジェクション、TRTLLM NVFP4 MoE パス（量子化、2 つのルーティングカーネル、2 つのグループ化 GEMM）。残りの時間の大部分はこれらのライブラリカーネルに費やされており、vLLM はこれらを次の最適化ターゲットと特定しています。

アテンション DP と MegaMoE で毎秒 168.8K トークンに到達

高並列領域では、アテンションデータ並列処理と MegaMoE を組み合わせることで Beam のパフォーマンスがさらに向上します。MegaMoE は FlashInfer の CuTe DSL NVFP4 MegaMoE カーネルを利用し、通信とエキスパート計算を融合できます。Finalize + ノルム境界融合はこの構成でも適用可能です。

構成 ピーク入力トークンスループット
アテンション DP + MegaMoE 168.8K トークン / 秒
TP (flashinfer_trtllm, flashinfer_nvlink_one_sided) 98.7K トークン / 秒

アテンション DP と MegaMoE は、TP4 よりもピークプリフィルスループットが 1.71 倍高くなりました。ベンチマークは 4 x GB300 を使用し、入力シーケンス長 32K トークン、出力シーケンス長 1 トークンでプリフィル性能を分離しました。

有効にするには、MoE と並列設定を次のように置き換えます：

--moe-backend flashinfer_moe_ep_cutedsl \
  --data-parallel-size 4 \
  --enable-expert-parallel

DFlash2 投機的デコーディング: 並列度 1 で毎秒 708 トークン

Beam には DFlash2 ドラフトヘッドが同梱されており、vLLM はデイ 0 からサポートしています。DFlash2 は一度に 1 トークンをドラフトする代わりに、ターゲットモデルの隠れ状態から単一のフォワードパスでトークンのブロック全体を予測します。

ドラフターは 5 つのスライディングウィンドウ層で構成され、5 つのターゲット層から隠れ状態を読み取り、ターゲットの埋め込みと LM ヘッドを共有します。グループ化畳み込みと小さな候補セレクターが、並列ドラフターが失う左から右への依存関係の一部を復元します。

適応的検証は、"enable_adaptive_verification": true で有効になり、各ドラフトトークンが受け入れられる可能性を予測し、ステップごとに固定数ではなく、コストに見合うものだけを検証します。DFlash2 には信頼度ヘッドがないため、受け入れ推定はドラフター自身のロジットから得られます。ドラフターの畳み込み、ノルム、KV 書き込みは融合され、ドラフトステップ全体が CUDA グラフで実行されます。FP8 KV キャッシュはドラフターでも機能し、受け入れ率は変わりません。

TP8 でのサービング結果は、SPEED-Bench の throughput_1k 分割の低エントロピーカテゴリからのものです：約 1K トークンのプロンプトで、主にコード補完とテキストソート、グリーディデコーディング、最大 4K 出力トークン。低エントロピーテキストは予測が容易なため、オープンエンドのチャットでは受け入れ率が低くなることが予想されます。

並列度 固定ドラフト (tok/s) 適応的 (tok/s) 変化 受容長
1 708 707 -0.1% 3.71 -> 3.61
8 3,232 3,221 -0.3% 3.67 -> 3.57
32 7,744 8,275 +6.9% 3.68 -> 3.29
128 15,012 16,632 +10.8% 3.68 -> 3.23

適応的検証は負荷が高いときに効果を発揮します。並列度 128 では出力スループットが 10.8% 向上し、p99 TPOT が 9% 削減されます。一方、ステップごとに検証するドラフトトークン数は減ります。受容長はわずかに低下しますが、節約される計算量がそれを上回ります。低並列度では、追加のドラフトトークンがほぼ無料であるため、効果は中立です。

デプロイメントガイダンス: 並列設定、推論、ツール呼び出し

vLLM の投稿の冒頭にあるコマンドが推奨デフォルトです：単一 GB300 ノードで TP4、DFlash2 と FP8 KV キャッシュを使用。適応的検証を有効にするには、--speculative-config に "enable_adaptive_verification": true を追加します。上記の他の最適化はデフォルトで有効です。

最低レイテンシには TP4 または TP8 を使用します。より高いスループットには、データ並列アテンションとエキスパート並列処理を使用します。例：-dp 4 --enable-expert-parallel。

beam 推論およびツール呼び出しパーサーは、チェックポイントに同梱されている応答テンプレートを読み取ります。デフォルトでは、vLLM はすべての応答をそのテンプレートに制約するため、ツール呼び出しは常に解析され、各ツールの JSON スキーマに一致します。response_format と thinking_token_budget は通常どおり機能します。デフォルトの制約をオフにするには、VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT=0 を設定します。

SkyRL によるデイ 0 強化学習サポート

vLLM は SkyRL チームと協力し、Beam での強化学習のデイ 0 サポートを可能にしました。LoRA、BF16 フルファインチューニング、MXFP8 RL がそのままサポートされ、500B スケール向けにアップグレードされた RL スタックが提供されます。

SkyRL チームは、vLLM を推論に、Megatron をトレーニングに使用して、DAPO レシピで Beam をポストトレーニングしました。LoRA、フルファインチューニング、MXFP8 RL は、50 トレーニングステップ（または 100 オプティマイザ更新、推論 minimal）で同様のパフォーマンス向上を達成し、AIME25 での平均精度（avg@8）が約 61% から約 79% に跳ね上がりました。Megatron と vLLM の間の平均絶対ログ確率差は、LoRA と BF16 フルファインチューニングで Rollout Routing Replay (R3) を有効にすると 0.01 未満、MXFP8 RL では 0.03 未満に保たれます。

Beam サービングスタックの開発者

vLLM はこの作業を Inferact、Reflection、Nvidia チーム間の緊密なコラボレーションと説明しています。Inferact の貢献者は Woosuk Kwon、Ziqi Zhao、Ao Shen、Tiezhen Wang です。Reflection の貢献者は Kaixing Wu、Yoni Gozlan、Neil Gollapudi、Thibaut Barroyer、Aymen Kallala、Matthias Reso、Florent Gbelidji です。Nvidia の貢献者は Siyuan Fu です。

Sources