Reflection Beam 1.0 在 vLLM 上:Day 0 服务、代理优化内核和 DFlash2 投机解码
vLLM 为 Reflection 的 Beam 增加了 Day 0 服务支持,这是一个 501B 参数的混合专家模型,每个 token 有 23B 活跃参数,并专门针对 NVIDIA Blackwell 硬件进行了优化。这一成果之所以重要,是因为它结合了完整的生产级服务功能(推理、工具调用、结构化输出、张量/数据/专家并行)与编码代理自主优化的内核栈,使得 batch-1 解码速度比 vLLM 的第一个可用实现快 2.71 倍,预填充速度快 1.87 倍。
在 TP8 上,使用 Beam 的 DFlash2 草稿头,单个请求可达到每秒 708 个 token,接受长度为 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 是一个 52 层的 MoE 解码器,隐藏大小为 5120。其 MoE 路由遵循 DeepSeek-V3;架构的独特之处在于注意力层和残差流的归一化方式。
注意力:滑动窗口层加 NoPE 全注意力层
Beam 将注意力分为两种机制。43 层使用滑动窗口注意力,覆盖最后 512 个 token,并采用 RoPE。其余 9 层(每第六层,加上最后一层)在无位置编码(NoPE)的情况下关注整个上下文。
由于在没有 RoPE 时,随着键数量的增加,注意力会变得平坦,Beam 在 NoPE 层上应用 SSMax。SSMax 将每个查询按 (1 + p/2048)^0.1 缩放,其中 p 是基于 1 的位置,因此随着上下文增长,softmax 会变得更加尖锐。
Beam 使用分组查询注意力,有 48 个查询头、8 个 KV 头和 128 的头维度。Q 和 K 在 RoPE 或 SSMax 之前按头进行 RMS 归一化。在输出投影(o_proj)之前,每个头的输出乘以 sigma(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 和残差流:FP32 路由器、三明治归一化、深度缩放残差
Beam 有 256 个路由专家,d_ff 为 2,560,使用 NVFP4。路由与 DeepSeek-V3 匹配:sigmoid 分数加上校正偏置选择 8 个专家组中最好的 4 个,然后选择前 8 个专家,重新归一化的权重乘以 2.5。与 DeepSeek-V3 不同,Beam 没有共享专家。第 0 到第 2 层使用密集的 SwiGLU MLP,d_ff 为 20,480,而不是 MoE。
路由器完全使用 FP32:5120 x 256 的路由器权重、累加和 logits 都在 FP32 中运行。
两个归一化选择塑造了残差流。每个子层前后都有一个 RMSNorm("三明治归一化")。第 n 次残差加法计算 h <- sqrt((n-1)/n) * h + x/sqrt(n),而不是 h + x,因此当子层输出具有单位方差时,残差流在每个深度保持单位方差,而不是随深度增长。
代理驱动的优化使 batch-1 解码速度提升 2.71 倍
vLLM 的第一个 Beam 实现未优化,几乎纯 PyTorch。随后编码代理自主优化了它:他们分析了模型,编写了内核,并对每个更改进行了基准测试。在 TP8 上,七项更改将 batch-1 解码从每步 8.95 毫秒降至 3.30 毫秒,8 x 2K 预填充从 328 毫秒降至 176 毫秒。
七项更改是:
- 融合归一化、残差和归一化(-40% 解码)。在 104 个子层边界处的后归一化、缩放残差加法和下一个前归一化作为一个 Triton 内核运行,使用 FP32 数学和参考实现的 BF16 舍入点。
- 融合 Q/K 归一化、RoPE/SSMax 和输出门。每头 Q/K RMSNorm 和 RoPE(或 SSMax)作为一个内核运行,sigmoid 门作为另一个内核。
- 将两个张量并行全归约融合到归一化内核中,基于 FlashInfer 的多播全归约。在 MoE 之后,同一内核还组合了前 8 个专家的输出。
- 为 FP32 路由器专门设计的 GEMM 内核。路由器之前回退到通用路径;vLLM 的小批量 FP32 路由器内核现在支持 Beam 的 5120 x 256 形状,同时保持 FP32 权重、累加和 logits。
- 使用低延迟 GEMM 合并 QKV 和门投影。Q、K、V 和门权重在加载时打包到一个矩阵中,因此四个投影作为一个 GEMM 运行;对于小批量,它使用低延迟 CuTe DSL GEMM 或 FlashInfer 的 TinyGEMM。
- 融合 KV 缓存更新。Q/K 归一化和 RoPE 内核也写入 KV 缓存,消除了单独的缓存写入内核。
- 三个密集层的密集 MLP 激活在一个内核中完成。
延迟测量是在 GB200 上以 TP8 模型前向延迟进行的,针对 batch-1 解码和大约 100K 上下文的 8 x 2,048 token 预填充。
融合解码图在 13 个内核中运行一个 MoE 层
经过这些更改,在 TP8 上的 batch-1 解码中,一个 MoE 层作为 13 个内核运行。其中六个内核是为 Beam 编写、融合或扩展的;每个小的 Beam 特定操作都合并到相邻操作中,两个全归约都位于归一化内核内部。
其他七个内核来自库:TRTLLM-gen 注意力、cuBLAS 输出投影和 TRTLLM NVFP4 MoE 路径(量化、两个路由内核和两个分组 GEMM)。大部分剩余时间花在这些库内核上,vLLM 将其确定为下一个优化目标。
注意力 DP 加 MegaMoE 达到每秒 168.8K token
在高并发区域,通过将注意力数据并行与 MegaMoE 结合,Beam 性能进一步提高,这得益于 FlashInfer 的 CuTe DSL NVFP4 MegaMoE 内核,它可以融合通信与专家计算。Finalize + 归一化边界融合在此配置中仍然适用。
| 配置 | 峰值输入 token 吞吐量 |
|---|---|
| 注意力 DP + MegaMoE | 168.8K token / 秒 |
TP(flashinfer_trtllm、flashinfer_nvlink_one_sided) |
98.7K token / 秒 |
注意力 DP 与 MegaMoE 的峰值预填充吞吐量比 TP4 高 1.71 倍。基准测试使用 4 x GB300,输入序列长度为 32K token,输出序列长度为 1 token,以隔离预填充性能。
要启用它,请将 MoE 和并行设置替换为:
--moe-backend flashinfer_moe_ep_cutedsl \
--data-parallel-size 4 \
--enable-expert-parallel
DFlash2 投机解码:并发 1 时每秒 708 token
Beam 附带一个 DFlash2 草稿头,vLLM 从 Day 0 开始支持。DFlash2 不是一次草拟一个 token,而是在一次前向传递中从目标模型的隐藏状态预测整个 token 块。
草稿器由 5 个滑动窗口层构建,这些层从 5 个目标层读取隐藏状态,并共享目标的嵌入和 LM 头。分组卷积和小型候选选择器恢复了一些并行草稿器丢失的从左到右的依赖性。
自适应验证,通过 "enable_adaptive_verification": true 启用,预测每个草稿 token 被接受的可能性,并只验证那些值得其成本的 token,而不是每步固定数量。DFlash2 没有置信度头,因此接受估计来自草稿器自身的 logits。草稿器的卷积、归一化和 KV 写入是融合的,整个草稿步骤在 CUDA 图中运行。FP8 KV 缓存也适用于草稿器,接受率没有变化。
TP8 上的服务结果来自 SPEED-Bench 的 throughput_1k 分区的低熵类别:大约 1K token 的提示,主要是代码补全和文本排序,使用贪婪解码和最多 4K 输出 token。低熵文本更容易预测,因此在开放式聊天中应预期较低的接受率。
| 并发 | 固定草稿(token/秒) | 自适应(token/秒) | 变化 | 接受长度 |
|---|---|---|---|---|
| 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 几乎免费,效果是中性的。
部署指南:并行设置、推理和工具调用
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 的 Day 0 强化学习支持
vLLM 与 SkyRL 团队合作,为 Beam 上的强化学习提供 Day 0 支持。LoRA、BF16 全微调和 MXFP8 RL 开箱即用,并升级了适用于 500B 规模的 RL 栈。
SkyRL 团队使用 vLLM 进行推理,使用 Megatron 进行训练,在 DAPO 配方上对 Beam 进行了后训练。LoRA、全微调和 MXFP8 RL 在 50 个训练步骤(或 100 个优化器更新,推理 minimal)中取得了类似的性能提升,在 AIME25 上平均准确率(avg@8)从约 61% 跃升至约 79%。在启用 Rollout Routing Replay (R3) 的情况下,Megatron 和 vLLM 之间的平均绝对对数概率差异保持在 0.01 以下(LoRA 和 BF16 全微调),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。