vLLM 从 Day 0 支持 Reflection Beam，并针对 NVIDIA Blackwell 优化

TL;DR：Reflection 的 Beam 是一个 501B 参数的 MoE，每个 token 激活 23B 参数。vLLM 从第 0 天起就支持它，并在 NVIDIA Blackwell 上对其进行了优化：

完整功能支持。推理、工具调用和结构化输出开箱即用，你可以用张量并行、数据并行或专家并行来运行该模型。
智能体驱动的优化。编码智能体自主优化了该模型，使 batch-1 解码比我们的首个实现快 2.71 倍，预填充快 1.87 倍。
带自适应验证的 DFlash2。使用 Beam 的 DFlash2 草稿模型，单个请求在 TP8 上达到 708 tokens/s，接受长度为 3.7。自适应验证会在每一步决定有多少草稿 token 值得验证，在高并发下额外带来 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 层的 MoE 解码器，隐藏维度为 5120。MoE 路由遵循 DeepSeek-V3。主要差异在于注意力层以及残差流的归一化方式。

图 1：Beam 的层堆叠和一个稀疏解码器层。金色标记的是与标准 pre-norm MoE 不同的部分。全屏打开。

注意力

混合滑动窗口和 NoPE 层。43 层对最后 512 个 token 使用带 RoPE 的滑动窗口注意力。其余 9 层（每第六层，加上最后一层）在完整上下文上做注意力，且不使用位置编码（NoPE）。
NoPE 层上的 SSMax。没有 RoPE 时，注意力会随着键数量增长而变得平坦。SSMax 将每个查询按 (1 + p/2048)^0.1 缩放，其中 p 是从 1 开始的位置，因此 softmax 会随着上下文增长而变得更锐利。
带逐头 Q/K 归一化的 GQA。48 个查询头，8 个 KV 头，头维度 128。Q 和 K 在 RoPE 或 SSMax 之前按头进行 RMS 归一化。
Sigmoid 输出门。在 o_proj 之前，每个头的输出乘以 σ(gate + bias)，其中 gate 来自其自身的全宽投影，并带有一个学习到的逐头偏置。一个头可以逐 token 被调低或关闭。
Small KV cache. Only the 9 full-attention layers keep KV for the whole sequence:
- 每个 token 9 × 2 × 8 × 128 个值，即 BF16 下 36 KiB，FP8 下 18 KiB。
- 在分片之前，完整 1M token 上下文约需 18 GiB 的 FP8 KV。

MoE 和残差流

DeepSeek-V3 风格 MoE。256 个路由专家（d_ff 2,560，NVFP4），路由方式与 DeepSeek-V3 相同：sigmoid 分数加上校正偏置选出 8 个专家组中最好的 4 个，然后选出前 8 个专家，重新归一化后的权重乘以 2.5。与 DeepSeek-V3 不同，这里没有共享专家。第 0–2 层改用稠密 SwiGLU MLP（d_ff 20,480）。
FP32 路由器。5120 × 256 的路由器权重、累加和 logits 全部为 FP32。
三明治归一化。每个子层之前有一个 RMSNorm，之后还有一个。
深度缩放残差。第 n 次残差加法计算 h ← √((n−1)/n)·h + x/√n 而不是 h + x。在子层输出为单位方差的情况下，残差流在每个深度都保持单位方差，而不是随之增长。

智能体驱动的优化

我们的初始实现版本是一个未优化、几乎纯 PyTorch 的模型版本。随后我们用编码智能体自主优化它：它们对模型进行性能分析、编写内核，并对每项改动进行基准测试。在 TP8 上，以下七项改动将 batch-1 解码从每步 8.95 ms 降至 3.30 ms，并将 8 × 2K 预填充从 328 ms 降至 176 ms。

图 2：每次变更后在 GB200 上的 TP8 模型前向延迟，针对 batch-1 解码和约 100K 上下文下的 8 × 2,048 token 预填充。两个面板使用不同的刻度。查看表格获取精确值，或全屏打开。

融合 norm → residual → norm（解码 −40%）。每个 104 个子层边界处的 post-norm、缩放残差加法和下一个 pre-norm 现在作为一个 Triton kernel 运行，使用 FP32 数学和参考实现的 BF16 舍入点。
融合 Q/K norm、RoPE/SSMax 和输出门。 每头的 Q/K RMSNorm 和 RoPE（或 SSMax）作为一个 kernel 运行，sigmoid 门作为另一个。
将 all-reduce 融合进 norm。 每层中的两个 TP all-reduce 都在 norm kernel 内运行，基于 FlashInfer 的 multicast all-reduce。在 MoE 之后，同一个 kernel 还会合并 top-8 专家输出。
为 FP32 router 定制的 GEMM kernel。 router 过去会回退到通用路径。vLLM 的小批量 FP32 router kernel 现在支持 Beam 的 5120 × 256 形状，保持 FP32 权重、累加和 logits。
将 QKV + gate 投影与低延迟 GEMM 合并。 Q、K、V 和 gate 权重在加载时打包成一个矩阵，因此四个投影作为一个 GEMM 运行。对于小批量，它使用低延迟 CuTe DSL GEMM 或 FlashInfer 的 TinyGEMM。
融合 KV-cache 更新。 Q/K norm 和 RoPE kernel 现在还会将 K 和 V 写入 KV cache，因此没有单独的 cache 写入 kernel。
在一个 kernel 中完成 Dense-MLP 激活，用于三个 dense 层。

融合图

在这些变更之后，一个 MoE 层在 TP8 上的 batch-1 解码中作为 13 个 kernel 运行。

图 3：一个 MoE 层按启动顺序排列的 kernel，来自 batch-1 TP8 解码的 rank-0 trace。金色 kernel 是为 Beam 编写、融合或扩展的。每个 kernel 列出它计算的操作。全屏打开。

为 Beam 编写、融合或扩展了 6 个 kernel。 每个小的 Beam 特定操作都被合并到相邻操作中，两个 all-reduce 都位于 norm kernel 内。
7 个 kernel 来自库： TRTLLM-gen attention、cuBLAS 输出投影，以及 TRTLLM NVFP4 MoE 路径（量化、两个路由 kernel、两个分组 GEMM）。剩余时间大部分在这里，所以它们是下一步。

Attention DP + MegaMoE

图 4：FlashInfer CuteDSL NVFP4 MegaMoE 可以将通信与专家计算融合。

在高并发区域，通过将 Attention DP 与 MegaMoE 结合，Beam 的性能可以进一步提升，这由 FlashInfer 的 CuTe DSL NVFP4 MegaMoE kernel 提供支持。Finalize + norm 边界融合仍然适用。

配置	峰值输入 token 吞吐量
Attention DP + MegaMoE	168.8K tokens / 秒
TP（`flashinfer_trtllm`、`flashinfer_nvlink_one_sided`）	98.7K tokens / 秒

Attention DP 配合 MegaMoE 的峰值预填充吞吐量比 TP4 高 1.71 倍。基准设置：4×GB300，输入序列长度为 32K token，输出序列长度为 1 token，以隔离预填充性能。

要启用 Attention DP + MegaMoE，请将 MoE 和并行设置替换为：

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

DFlash2 投机解码

Beam 附带一个 DFlash2 草稿头，vLLM 从第 0 天起就支持。它不是一次草拟一个 token，而是从目标模型的隐藏状态中通过单次前向传播预测整个 token 块。

Drafter。 5 个滑动窗口层，从 5 个目标层读取隐藏状态，并共享目标的嵌入和 LM head。分组卷积和一个小型候选选择器恢复了并行 drafter 丢失的部分从左到右依赖。
自适应验证。 借助 "enable_adaptive_verification": true，vLLM 预测每个草稿 token 被接受的可能性，并只验证值得其成本的那些，而不是每步固定数量。DFlash2 没有置信度头，因此接受估计来自 drafter 自身的 logits（参见我们关于自适应验证的近期文章）。
内核。 drafter 的卷积、归一化和 KV 写入被融合，整个草稿步骤在 CUDA 图中运行。FP8 KV 缓存也适用于 drafter，且接受率不变。

在 TP8 上的服务结果，包含和不包含自适应验证，在 SPEED-Bench 的 throughput_1k 分割的低熵类别上：约 1K token 的提示，主要是代码补全和文本排序，使用贪心解码，输出最多 4K token。低熵文本更容易预测，因此在开放式聊天中预期接受率更低。

并发数	固定草稿 (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%，同时每步验证更少的草稿 token，这略微降低了接受长度，但节省的计算量超过损失。在低并发下，额外草稿 token 几乎免费，因此它是中性的。

快速开始

本文顶部的命令是一个不错的默认值：在单个 GB300 节点上使用 TP4，搭配 DFlash2 和 FP8 KV 缓存。要启用自适应验证，请将 "enable_adaptive_verification": true 添加到 --speculative-config。上述其他优化默认开启。

Choosing a parallel setup.
- 对于最低延迟，使用 TP4 或 TP8。
- 对于更高吞吐量，使用数据并行注意力与专家并行，例如 -dp 4 --enable-expert-parallel。
推理和工具调用。 beam 解析器读取检查点附带的响应模板。默认情况下，vLLM 还将每个响应约束到该模板，因此工具调用始终能解析并匹配每个工具的 JSON schema。response_format 和 thinking_token_budget 照常工作。要关闭默认约束，请设置 VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT=0。

SkyRL RL 的 Day 0 支持

我们还与 SkyRL 团队合作，为 Beam 上的 RL 提供 Day 0 支持。LoRA、BF16 全量微调和 MXFP8 RL 开箱即用，并针对 500B 规模升级了 RL 栈。

SkyRL 团队使用 DAPO 配方对 Beam 模型进行了后训练，推理使用 vLLM，训练使用 Megatron。LoRA、全量微调和 MXFP8 RL 在 50 个训练步骤（即 100 次优化器更新，推理 minimal）中实现了相似的性能提升，在 AIME25 上的平均准确率（avg@8）从约 61% 跃升至约 79%。在 LoRA 和 BF16 全量微调启用 Rollout Routing Replay (R3) 的情况下，Megatron 和 vLLM 之间的平均绝对对数概率差异保持在 0.01 以下，MXFP8 RL 则保持在 0.03 以下。

查看 SkyRL 博客了解更多详情！

致谢

这项工作由 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