vLLM 实现 Reflection Beam 501B MoE 模型 Day 0 支持

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 层使用带 RoPE 的滑动窗口注意力，窗口为最近 512 个 token。其余 9 层（每第六层，加上最后一层）在完整上下文上做注意力，且不使用位置编码（NoPE）。
• NoPE 层上的 SSMax。没有 RoPE 时，注意力会随着键数量增长而变得平坦。SSMax 将每个 query 乘以 (1 + p/2048)^0.1，其中 p 是从 1 开始的位置，因此 softmax 会随着上下文增长而变得更锐利。
• 带逐头 Q/K 归一化的 GQA。48 个 query 头，8 个 KV 头，头维度 128。Q 和 K 在 RoPE 或 SSMax 之前按头进行 RMS 归一化。
• Sigmoid 输出门。在 o_proj 之前，每个头的输出乘以 σ(gate + bias)，其中 gate 来自其自身的全宽投影，并带有一个学习到的逐头 bias。一个头可以逐 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 分数加校正 bias 从 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 的模型版本。随后我们用编码智能体自主优化它：它们对模型进行性能分析、编写 kernel，并对每项改动进行基准测试。在 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 kernel 运行，使用 FP32 数学和参考实现的 BF16 舍入点。
2. 融合 Q/K norm、RoPE/SSMax 和输出门。 每头的 Q/K RMSNorm 和 RoPE（或 SSMax）作为一个 kernel 运行，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-write kernel。
7. 在一个 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 可以将通信与专家计算融合。

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

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

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

自适应验证在负载下带来收益。在并发数 128 时，它将输出吞吐量提高 10.8%，并将 p99 TPOT 降低 9%，同时每步验证更少的 draft token，这略微降低了接受长度，但节省的计算量超过损失。在低并发下，额外 draft token 几乎免费，因此它是中性的。

快速开始

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

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

SkyRL 强化学习的 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 之间的平均绝对 log 概率差保持在 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