vLLM 首日支持 Reflection Beam 501B MoE 模型

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

• 完整功能支持。推理、工具调用和结构化输出开箱即用，你可以用张量并行、数据并行或专家并行来运行该模型。
• 智能体驱动的优化。编码智能体自主优化了该模型，使 batch-1 解码比我们的首个实现快 2.71 倍，预填充快 1.87 倍。
• 带自适应验证的 DFlash2。借助 Beam 的 DFlash2 草稿模型，单个请求可达到708 tokens/s在 TP8 上，接受长度为 3.7。自适应验证会在每一步决定有多少草稿 token 值得验证，从而在高并发下额外带来10% 的吞吐提升。

要试用它，请使用我们的 Docker 镜像vllm/vllm-openai:reflection-beam：

dockerrun--gpusall--ipc=host-p8000:8000\
-v~/.cache/huggingface:/root/.cache/huggingface\
vllm/vllm-openai:reflection-beam\
reflection/Beam-1.0-501B-A23B-NVFP4\
-tp4\
--max-num-seqs128\
--reasoning-parserbeam\
--tool-call-parserbeam--enable-auto-tool-choice\
--kv-cache-dtypefp8\
--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 调低或关闭。
• 小型 KV 缓存。只有 9 个全注意力层为整个序列保留 KV：
  • 每个 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 预填充。两个面板使用不同的刻度。查看表格获取精确数值，或全屏打开。

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

融合后的计算图

经过这些改动后，一个 MoE 层在13 个 kernel中运行（TP8 上的 batch-1 decode）。

图 3：一个 MoE 层按启动顺序排列的 kernel，来自 batch-1 TP8 decode 的 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 可以将通信与专家计算融合。

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

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

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

--moe-backendflashinfer_moe_ep_cutedsl\
--data-parallel-size4\
--enable-expert-parallel

DFlash2 投机解码

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

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

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

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

快速开始

本文顶部的命令是一个不错的默认配置：在单个 GB300 节点上使用 TP4，并启用 DFlash2 和 FP8 KV 缓存。要启用自适应验证，请添加"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。

Day 0 支持使用 SkyRL 进行 RL

我们还与 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