从第 0 天起在 vLLM 上提供 Reflection Beam 服务

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 在 TP8 上，接受长度为 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)，其中门控来自其自身的全宽度投影，并带有可学习的逐头偏置。可以逐 token 地将某个头调低或关闭。
• 小的 KV 缓存。 只有 9 个全注意力层为整个序列保留 KV：
  • 每个 token 9 × 2 × 8 × 128 个值，在 BF16 下为 36 KiB，在 FP8 下为 18 KiB。
  • 对于完整的 1M-token 上下文，在分片之前，FP8 KV 约需 18 GiB。

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，之后还有一个 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 预填充。两个面板使用不同的比例。查看表格以获取精确数值，或 全屏打开它.

1. 融合的归一化 → 残差 → 归一化 （解码减少 40%）。在 104 个子层边界处，每个后置归一化、带缩放的残差相加以及下一个前置归一化现在作为一个 Triton 内核运行，使用 FP32 数学运算并保持参考实现的 BF16 舍入点。
2. 融合的 Q/K 归一化、RoPE/SSMax 和输出门控。 逐头的 Q/K RMSNorm 和 RoPE（或 SSMax）作为一个内核运行，sigmoid 门控作为另一个内核运行。
3. All-reduce 融合进归一化内核。 每层中的两个 TP all-reduce 都在归一化内核内运行，该内核基于 FlashInfer 的多播 all-reduce 构建。在 MoE 之后，同一个内核还合并前 8 个专家的输出。
4. 为 FP32 路由器定制的 GEMM 内核。 路由器此前会回退到通用路径。vLLM 的小批量 FP32 路由器内核现在支持 Beam 的 5120 × 256 形状，保持 FP32 权重、累加和 logits。
5. 使用低延迟 GEMM 合并 QKV + 门控投影。 Q、K、V 和门控权重在加载时打包成一个矩阵，因此这四个投影作为一次 GEMM 运行。对于小批量，它使用低延迟的 CuTe DSL GEMM 或 FlashInfer 的 TinyGEMM。
6. 融合的 KV-cache 更新。 Q/K norm 和 RoPE 内核现在还会将 K 和 V 写入 KV cache，因此不再需要单独的 cache 写入内核。
7. 单内核中的 Dense-MLP 激活 用于三个 dense 层。

融合后的计算图

在这些更改之后，一个 MoE 层运行为 13 个内核 在 TP8 上 batch-1 解码时。

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

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

Attention DP + MegaMoE

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

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

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

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

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

DFlash2 投机解码

Beam 自带一个 DFlash2 draft head，vLLM 从第 0 天起就支持它。它不是一次起草一个 token，而是在单次前向传播中从目标模型的 hidden states 预测整个 token 块。

• 起草器。 5 个 sliding-window 层，从 5 个目标层读取 hidden states，并共享目标的 embedding 和 LM head。一个 grouped convolution 和一个小型候选选择器恢复了并行起草器所损失的 部分 从左到右的依赖关系。
• 自适应验证。 借助 "enable_adaptive_verification": true，vLLM 会预测每个草稿 token 被接受的可能性，只验证值得花费成本的那些，而不是每步固定验证一定数量。DFlash2 没有置信度头，因此接受率估计来自起草器自身的 logits（参见我们的 关于自适应验证的最新文章).
• 内核。 起草器的卷积、归一化和 KV 写入被融合在一起，整个草稿步骤在 CUDA 图中运行。FP8 KV cache 也适用于起草器，且接受率没有变化。

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

自适应验证在高负载下才见成效。在并发 128 时，它将输出吞吐量提升 10.8%，并将 p99 TPOT 降低 9%，同时每步验证的草稿 token 更少，这使接受长度略有下降，但节省的计算量超过了损失。在低并发时，额外的草稿 token 几乎没有成本，因此它是中性的。

开始使用

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

• 选择并行安装方式。
  • 要获得最低延迟，请使用 TP4 或 TP8。
  • 为了获得更高的吞吐量，请使用结合专家并行的数据并行注意力，例如： -dp 4 --enable-expert-parallel.
• 推理和工具调用。 该 beam 解析器读取检查点自带的响应模板。默认情况下，vLLM 还会将每个响应约束到该模板，因此工具调用始终能够被解析并匹配每个工具的 JSON schema。 response_format 和 thinking_token_budget 照常工作。要关闭默认约束，请设置 VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT=0.

SkyRL 强化学习第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
• 反思： 吴凯兴、约尼·戈兹兰、尼尔·戈拉普迪、蒂博·巴罗耶、艾曼·卡拉拉、马蒂亚斯·雷索、弗洛朗·贝利吉
• Nvidia: Siyuan Fu