跳到正文
LMSYS Blog· LMSYS Blog·· 17 小时前AI 评分62

SGLang 原生支持 Reflection 的 Beam 并驱动其生产 API

AI 导读

SGLang 宣布原生支持 Reflection 的 Beam 模型，并为其生产 API 提供支持。Beam 是 501B 参数的 MoE 模型，每个 token 激活 23B 参数，原生 1M token 上下文窗口，计划以 Apache 2.0 许可发布。

正文

Reflection 的 Beam 是一个 501B 参数的混合专家模型，每个 token 激活 23B 参数，专为软件工程和长周期智能体任务打造。Beam 计划以 Apache 2.0 许可发布，旨在让开发者能够在他们已经在使用的运行时、框架和编码工具链上部署和构建。

SGLang 现已原生支持 Beam，并为 Reflection 的生产 API 提供支持。它帮助开发者将 Beam 的推理效率转化为经过生产验证的实用服务性能。BF16、NVFP4 和动态 FP8 支持在精度和内存需求方面提供了灵活性；DFlash2 投机解码加速生成；前缀缓存与混合注意力内存管理相结合，实现了智能体交互间上下文的高效复用。在这些服务能力之下，我们优化了 Beam 的投影、归一化和通信，以提升吞吐量。

本文介绍 Beam 在 SGLang 上的集成、优化和性能结果。要部署该模型，请从我们的 Beam 服务指南 开始，获取经过验证的配置和启动说明。

亮点

  • Beam 1.0 501B-A23B： 一个仅解码器 MoE，采用滑动窗口和全注意力，256 个路由专家和 top-8 选择，以及原生 1M token 上下文窗口。
  • 检查点支持： BF16、ModelOpt NVFP4 和 compressed-tensors 动态 FP8。
  • 服务特性： DFlash2 投机解码、基数缓存和 HiCache、分块预填充、CUDA graphs、TP/EP/PP 和注意力数据并行。
  • 内核优化： 投影打包、跨注意力和残差边界的融合、用于 FP32 路由器的专用 BF16x3 GEMM，以及 GB300 MoE 调优。
  • 性能： 在 4 块 NVIDIA GB300 GPU 上使用 NVFP4 检查点进行 OpenHands 多轮智能体编码回放时，SGLang 在并发 32 下达到每 GPU 127.6K 总 token/秒，每用户 185.87 token/秒（p50）。

模型架构

Beam 将分组查询注意力与滑动窗口和全注意力层的混合结构相结合。其前馈堆栈从密集 SwiGLU 层开始，然后切换到路由 MoE。Beam 在投影前对注意力输出进行门控，并在两个分支周围放置三明治归一化。深度缩放残差控制每个分支如何加入隐藏状态流。

Beam's embedding and decoder stack, with an expanded decoder layer showing attention, feed-forward work, sandwich normalization and depth-scaled residual paths.

Beam 的解码器以及单层内的操作。注意力和前馈分支都以输出归一化和深度缩放残差加法结束。

组件原生配置
解码器层52
隐藏维度5,120
注意力头48 个查询头和 8 个 KV 头
头维度128
滑动窗口512 个 token
前馈调度三个初始密集层，随后是路由 MoE 层
路由专家256 个专家，每个 token 选择 8 个
上下文限制1,048,576 个 token

检查点指定了注意力和前馈调度。原生注意力调度在滑动窗口层组之后放置全注意力，并以全注意力结束。这为服务栈提供了两种需要管理的 KV 生命周期：全注意力的完整历史，以及滑动注意力的近期窗口。

混合注意力和输出门控

GQA 在 48 个查询头之间共享八个 KV 头，从而减少 SGLang 为每个 token 存储的 KV 缓存。滑动窗口层也将其活跃历史限制在最近的窗口内，并使用 RoPE。全注意力层使用 NoPE，并采用位置相关的 SSMax 查询缩放。两者都在注意力之前对 Q 和 K 进行归一化。

Beam 在输出投影之前对注意力结果进行门控。第二个投影从归一化后的隐藏状态生成门激活；模型添加一个学习到的逐头偏置，应用 sigmoid，并将注意力输出与逐元素门相乘。QKV 和门投影共享一个输入，这让我们可以在符合条件的布局中将它们打包成一个 BF16 投影。

稠密层与分组 MoE 路由

前三个前馈层使用稠密 SwiGLU。后续层从 256 个专家中选出八个专家，并组合它们的 SwiGLU 输出。

路由器计算 FP32 logits。分组 top-k 选择使用 sigmoid 分数和专家校正偏置来选择专家。后端保留 FP32 路由系数，并在专家组合期间应用 Beam 的路由缩放因子，或者在该内核需要时将其携带在路由权重中。

SSMax

Beam 的全注意力层使用可扩展 Softmax（SSMax），以在上下文增长时保持注意力聚焦。在 Q/K 归一化之后，位于从零开始的位置 p 的每个查询按 (1+(p+1)/2048)0.1 缩放。这会逐渐锐化注意力分布，抵消 softmax 权重在长上下文中倾向于分散到更多 token 的趋势。

三明治归一化与深度缩放残差

每个注意力和前馈分支都有输入和输出 RMS 归一化。Beam 按从层深度推导出的系数缩放分支和残差，然后将它们相加。两个缩放后的操作数在相加之前都会舍入到激活数据类型。

未融合的过渡会写入归一化后的分支，读取它进行残差组合，并读取组合结果用于下一个输入归一化。这些内存传递在解码器中重复进行，因此我们在保持 Beam 舍入边界的同时融合该过渡。最终解码器归一化和语言模型头馈入 FP32 logit softcapping。

功能支持

带数据并行注意力的 Dflash2 推测解码

Beam 通过 SGLang 的 DFLASH 运行时支持 DFlash2。目标模型在最终归一化之前捕获所选解码器层的层后残差。草稿模型使用这些特征来提议一个 token 块，Beam 在一次目标前向中验证该块。

Beam 的集成添加了一条新的 CUDA 服务路径，将 DFlash2 与注意力数据并行相结合。借助 --attn-dp-size N 和 --enable-dp-lm-head，运行时使注意力和草稿头工作与 DP 布局保持一致。SGLang 在 DP 填充确定最终 token 形状之后初始化目标验证元数据，包括存在空闲对等方的批次。

我们将草稿 KV 内存纳入混合池预算，以便目标缓存分配为草稿模型留出空间。草稿捕获要求 PP=1，并且服务上下文必须同时容纳目标模型和所选的草稿模型。

目标与草稿 KV 的统一内存

Beam 的全注意力 KV 随上下文增长，而其滑动层保留一个近期窗口。上游 SGLang 的统一 KV 分配器用共享的 GPU 页预算为这些分配提供支持。请求表保持稳定的虚拟 token ID；注意力读取和 KV 写入将它们转换为当前物理槽位。转换接口在两个方向上使用相同的虚拟到物理契约。

DFlash 进入统一池将此布局扩展到符合条件的全注意力 drafter。目标模型和 draft 模型在目标的页包络内保持不同的 KV 行，因此分配器将两者一起计入。紧凑窗口 drafter 可以基于目标的虚拟 ID 保留自己的请求到 token 表，限制其读取哪些页，同时其 KV 存储保持在共享布局中。Draft 写入在存储 KV 之前转换这些 ID。

放置取决于 draft 几何结构、KV dtype 和注意力后端。该门控在 HiCache、主机池回收或解码上下文并行的情况下拒绝共享 draft 放置。混合 SWA 主机会为私有 draft 回退预留预算；不支持的主机/draft 组合在启动时被拒绝。

我们将此路径扩展到具有解码 CUDA 图的 trtllm_mha 目标和 drafter。它使用为当前迭代读/写范围规划的页表，并将新页包络清零。在 SM90/SM120 上，它将 fmha_v2 预填充工作区与 XQA 解码信号量分离。该扩展在 TRTLLM 后端中通过目标验证和 draft 执行传递转换后的 KV 位置。

用于预填充/解码分离的异构 TP

预填充和解码处理不同的 token 批次形状，因此分离式服务器可以为每个阶段选择单独的注意力 TP 大小。对于 Beam 的混合注意力，当 KV 头所有权在两个组之间变化时，这种交接必须保留全注意力 KV 和活动滑动窗口。

PR #40909 通过 NIXL 和 Mooncake 为混合滑动窗口模型添加了此传输路径。它使用现有的 rank 所有权映射，以 token/头段复制窗口状态，将每个 token 的头段放入匹配的解码缓冲区。暂存携带全注意力 KV；窗口状态遵循头切片路径。

推理、工具调用和 LoRA

SGLang 读取 Beam 的检查点响应模板，为流式和非流式请求分离推理、答案和工具调用。该集成涵盖 Chat Completions、Responses 和 Messages，并支持 Beam 工具语法和 JSON 输出约束。

BF16 LoRA 服务使用 Triton MoE 运行器，涵盖注意力投影和门控、密集 MLP、路由专家、嵌入和输出头。加载器保留注意力门控的限定名称，以免与 MLP 门控目标冲突。

BF16、NVFP4 和 FP8 模型支持

我们支持原生 BF16、FP8 和 NVFP4 检查点。我们通过 FlashInfer MegaMoE 添加了 NVFP4 专家并行替代方案。它需要 flashinfer_megamoe 运行器和全对全后端；BF16 和动态 FP8 仍在此 Beam 扩展之外。优化部分解释了其 token 所有权和通信路径。

路由器在这些格式之间保持其精度和路由缩放契约。量化检查点和 LoRA 为其加载和适配器路径保留单独的 QKV 和注意力门控模块。

前缀复用、分块预填充和 HiCache

缓存的 Beam 前缀包含完整注意力历史以及在其复用边界处所需的滑动窗口状态。SGLang 将两者与分页缓存对齐，使请求能够从该前缀恢复，并保持其活动窗口完整。分块预填充将长提示分成多个部分，与解码工作一同处理。

HiCache 将前缀复用从 GPU 内存扩展到主机内存及受支持的存储后端。Beam 在启用 HiCache 的情况下保留其混合 SWA 池和滑动窗口内存节省。主机回载使用每个流水线阶段的本地层映射。

优化

融合优化

Beam 的归一化和门控路径围绕投影和集合通信串联起一系列小操作。我们融合这些链以消除中间存储、拷贝和启动开销。

  • 跨投影、归一化和集合通信边界的核融合。 我们融合 all-reduce、输出归一化、缩放残差加法和下一输入归一化，将可融合边界扩展到 NVFP4。PDL 让后续归一化在其依赖等待之前加载权重。滑动层将 Q/K 归一化与 RoPE 融合。完整注意力层使用融合的 Q/K 归一化并复用 SSMax 位置缩放。符合条件的 BF16 层将 QKV 和门控投影打包为一次启动。注意力门控融合偏置加法、sigmoid 和乘法。BF16 MoE 集合通信路径还集成了本地专家求和、路由缩放和 BF16 舍入。

  • 使用 MegaMoE 结合专家通信与计算。 我们在 FlashInfer MegaMoE 上启用 Beam 的 NVFP4 专家，将其 runner 与 MegaMoE all-to-all 后端配对。FlashInfer 在其融合 EP 核内执行专家分发、计算和合并。Beam 提供显式专家 ID 和路由权重，并接收每个 rank 令牌的完整输出。融合转换在加法之前保留两个缩放残差操作数的舍入。注意力和 MoE 有各自独立的融合资格检查；MoE 归约遵循所选后端的所有权。CPU 卸载、空流水线阶段和不支持的集合通信形状保留常规路径。所有 EP rank 都参与专家执行，包括没有本地令牌的 rank。

Token ownership with attention data parallelism and MegaMoE: dense MLP layers gather and scatter tokens, while sparse layers return complete outputs to each owning rank.

MegaMoE 令牌所有权与数据并行注意力。稠密 MLP 层保留 TP gather/scatter；稀疏层通过融合的专家分发、计算和合并传递 DP 本地令牌。完整输出返回到每个令牌所属的 rank，无需额外的 all-reduce。

BF16 MoE 核优化

预填充和解码为专家核提供不同的令牌形状。我们为 gate/up 和 down 投影添加了 GB300 tile 配置，并为选定的预填充形状使用 TMA。Triton 调优器还从其配置中读取 Beam 的专家维度。

对于单令牌解码，一条可选路径使用两个 GEMV 核计算路由专家输出。它保留 gate/up 和激活的 BF16 舍入边界，并在对 down 投影输出舍入之前应用 FP32 路由权重。调用方仍执行专家归约、路由缩放和 TP 通信。FP32 归约顺序与张量核路径不同。

TPU 支持

我们通过 SGLang-TorchTPU 为 Beam 带来首日 TPU 支持，这是一个用于在 TPU 上进行 LLM、VLM 和扩散模型推理的原生 PyTorch 后端。它将 TPU 加速与 SGLang 功能（如 RadixAttention 和 HiCache）相结合，为智能体工作负载实现经济高效、生产就绪的服务。

Beam 的 TPU 配方以及 SGLang-TorchTPU 的访问权限即将推出。通过此 Google Form 注册以获取最新动态。

Miles 中的强化学习

Miles 通过 Megatron 插件支持 Beam，并在 Reflection 的 SGLang 构建上运行 rollout。 训练后端实现了 Beam 的三明治归一化与深度缩放残差、 SSMax 注意力缩放、混合滑动窗口注意力、门控注意力输出、 无辅助损失的 MoE 路由，以及 FP32 最终 logit 软上限，并通过 Megatron-Bridge 直接加载 Hugging Face 检查点。目标是让同一响应的训练器与 rollout 对数概率尽可能接近。

并行性。 后端支持 DP、TP、SP、EP、专家 TP 和 PP。TP 对 注意力投影（包括输出门）进行分片，SP 则沿序列维度对它们之间的激活进行分片。EP 和专家 TP 独立于 注意力布局对路由专家进行分片，流水线阶段可以容纳不等数量的层。启用路由重放后， 训练器–rollout 差距在这些布局下保持同一水平， 且不会随训练步数增长。

路由重放。 Rollout Routing Replay 作为可选项提供。SGLang 记录每个生成 token 所选的专家，训练器的 MoE 层复用这些分配，而不是重新计算 top-k，因此近乎并列的路由器分数不会在训练中选择不同的专家路径。重放的路由在激活重计算期间以及每个流水线阶段也会生效。

数值一致性。 在 Beam 对精度敏感的地方，训练器遵循 SGLang 的数值处理。最终 logit 软上限以 FP32 运行，深度缩放残差加法在与引擎融合内核相同的点进行舍入。路由器门在训练、检查点和权重同步过程中始终保持 FP32。这些改动共同将训练器–rollout 对数概率差距减半，使其处于 BF16 归约顺序噪声的水平。确定性设置使重复运行可复现。

共置与分离部署。 Beam 可在任一模式下运行。分离式运行将训练和 rollout 放在不同的 GPU 上，并将下一次 rollout 与当前训练步重叠；TIS 缓解由此产生的一步策略陈旧所导致的离策略不匹配。共置运行共享相同的 GPU 并同步交替进行，在引擎生成时将训练器状态卸载到节点本地 NVMe。在两种模式下，更新后的 BF16 权重在每一步之后以桶的形式传输到引擎。

LoRA 训练支持。 LoRA 训练通过 Megatron-Bridge PEFT 进行，涵盖注意力、稠密 MLP 和路由专家投影。适配器以 Hugging Face 格式导出，并由 SGLang 直接提供服务。在 RL 期间，仅适配器权重通过权重更新会话流式传输到引擎，而基础权重保持常驻，引擎会验证每个接收到的适配器的校验和。

已验证运行

下图展示了在 DAPO-Math-17K 上进行的全参数 GRPO 运行，从早期的 Beam checkpoint 开始，在 96 块 H200 GPU 上以分离式异步模式运行。64 块 GPU 以 TP8、SP、EP8、expert TP8 和 DP8 进行训练，32 块 GPU 服务两个 TP16 SGLang 引擎。每次 rollout 生成 32 个 prompt × 8 个响应，响应上限为 4K token，并馈入一次优化器步骤。在前 204 步中，AIME25 准确率从 46.3% 上升到 56.7%，同时平均响应长度下降。运行期间，训练从保存的 checkpoint 恢复了一次。

Training dynamics: train-rollout KL divergence and gradient norm.

AIME 2025 accuracy and mean output length during training.

性能结果

OpenHands Agentic Coding 扫描

我们收集了 Beam NVFP4 在四块 GB300 GPU 上以 TP4 运行 OpenHands 多轮 agentic coding 工作负载的性能结果。每段对话以约 74K token 的 prompt 开始，每轮输出约 220 token，总共 13 轮。连续轮次复用前缀，聚合缓存命中率约 92%。我们扫描并发数 1、4、8、16 和 32 段对话，并绘制每 GPU token 吞吐量（tok/s/GPU）与 p50 交互性（tok/s/user）的关系。总吞吐量包括缓存输入 token。这些测量使用八 token 的 DFlash2 块，模拟接受长度固定为四，并使用真实 draft token；它们不测量自然 draft 接受率或模型准确率。

Beam NVFP4 OpenHands performance on four GB300 GPUs with TP4/EP1: total throughput per GPU versus p50 interactivity.

AgentX Agentic Coding 扫描

我们与 NVIDIA 合作，收集了 Beam NVFP4 在四块 GB300 GPU 上以 TP4/EP1 运行 AgentX 的性能结果，AgentX 是 SemiAnalysis 用于回放记录的多轮 agentic coding 轨迹的工具。AgentX 复现真实 coding-agent 会话的 prompt 大小、前缀共享、子代理扇出和轮次间时序，因此结果反映了“真实的 agentic 负载”。会话中的请求共享长前缀，因此 92% 到 96% 的输入 token 命中前缀缓存。使用了 FP8 KV 缓存，选定的投机块大小为 8。我们扫描并发数从 1 到 96 个会话，并绘制每 GPU token 吞吐量（tok/s/GPU）与 p90 交互性（tok/s/user）的关系。在并发数 96 时，吞吐量达到 83,504 tok/s/GPU，交互性为 43.9 tok/s/user；单个会话以 348.0 tok/s/user 运行。

Beam NVFP4 AgentX performance on GB300: total throughput per GPU versus p90 interactivity.

致谢

我们想向以下组织和个人表示感谢，感谢他们对 Beam 的 SGL 集成开发所做的贡献，以及对其性能进行的广泛基准测试。

SGLang 社区/RadixArk：Jimmy Shong、Andy Ye、Xiaoyu Zhang、Yanbin Jiang、Qiaolin Yu、Alex Nails、Yuhao Yang

Reflection：Aksel Joonas Reedi、Aymen Kallala、Christopher Vanderschuere、Ewen Wang、Florent Gbelidji、George Antoniadis、Joel Gustafson、Yoni Gozlan、Iaroslav Geraskin、Kai Wu、Neil Gollapudi、Renault Young、Asfandyar Qureshi、Thibaut Barroyer、Matthias Reso。

NVIDIA：Jinze Xue、Dong Xu、Michael Loh、Matthew Radzihovsky

附录

OpenHands 复现

Beam OpenHands 扫描背后的脚本位于公开 bench-openhands-gb300-sweep-beam 分支上的 benchmark/beam_openhands 中。从准备好的 Beam docker 镜像开始，然后使用 reproduce.py 安装隔离客户端、构建数据集，并在一个 4 × NVIDIA GB300 节点上以 TP4/EP1 运行并发级别 1、4、8、16 和 32。

服务器是该分支的 SGLang。客户端是提交 acd09b44 的 EvalScope，使用其 swe_smith 多轮工作负载和 Beam 回放适配器。HiCache 在并发数 16 和 32 时配置为每 rank 88 GB，在较低并发数时禁用。

该工作负载是基于 nebius/SWE-rebench-openhands-trajectories 构建的 OpenHands 多轮智能体回放，并用 nvidia/OpenScienceReasoning-2 文本进行填充：每段对话 13 轮，首轮 74,160 个裸输入 token，之后每轮额外增加 753 个输入 token，每轮输出 220 个 token。不断增长的对话历史使每个请求的 prompt token 约为 80K，聚合前缀缓存命中率达到 92%。

服务器启动命令如下，其中 <cc> 同时设置最大运行请求数和解码 CUDA 图批处理上限

export MODEL_PATH=/models/Beam-1.0-501B-A23B-NVFP4
export DRAFT_PATH=/models/Beam-1.0-501B-A23B-DFlash2
export SGLANG_SIMULATE_ACC_LEN=4
export SGLANG_SIMULATE_ACC_METHOD=match-expected
export SGLANG_SIMULATE_ACC_TOKEN_MODE=real-draft-token
export SGLANG_RESPONSE_TEMPLATE_TURN_GRAMMAR=0

export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
export PYTHONNOUSERSITE=1 TZ=UTC
export HF_HUB_OFFLINE=1

CC=16  # Set to 16 or 32
python3 -m sglang.launch_server \
    --model-path "$MODEL_PATH" \
    --served-model-name beam \
    --quantization modelopt_fp4 \
    --tp-size 4 \
    --ep-size 1 \
    --attn-dp-size 1 \
    --nnodes 1 \
    --node-rank 0 \
    --dist-init-addr 127.0.0.1:39746 \
    --attention-backend trtllm_mha \
    --page-size 64 \
    --mem-fraction-static 0.80 \
    --enable-unified-memory \
    --reasoning-parser beam \
    --tool-call-parser beam \
    --speculative-algorithm DFLASH \
    --speculative-draft-model-path "$DRAFT_PATH" \
    --speculative-draft-attention-backend trtllm_mha \
    --speculative-num-draft-tokens 8 \
    --enable-hierarchical-cache \
    --hicache-size 88 \
    --max-running-requests "$CC" \
    --cuda-graph-max-bs-decode "$CC" \
    --context-length 1048576 \
    --chunked-prefill-size 16384 \
    --prefill-decode-interval 0 \
    --enable-cache-report \
    --enable-metrics \
    --tokenizer-worker-num 1 \
    --host 0.0.0.0 \
    --port 30080

CC=1  # Set to 1, 4, or 8
python3 -m sglang.launch_server \
    --model-path "$MODEL_PATH" \
    --served-model-name beam \
    --quantization modelopt_fp4 \
    --tp-size 4 \
    --ep-size 1 \
    --attn-dp-size 1 \
    --nnodes 1 \
    --node-rank 0 \
    --dist-init-addr 127.0.0.1:39746 \
    --attention-backend trtllm_mha \
    --page-size 64 \
    --mem-fraction-static 0.80 \
    --enable-unified-memory \
    --reasoning-parser beam \
    --tool-call-parser beam \
    --speculative-algorithm DFLASH \
    --speculative-draft-model-path "$DRAFT_PATH" \
    --speculative-draft-attention-backend trtllm_mha \
    --speculative-num-draft-tokens 8 \
    --max-running-requests "$CC" \
    --cuda-graph-max-bs-decode "$CC" \
    --context-length 1048576 \
    --chunked-prefill-size 16384 \
    --prefill-decode-interval 0 \
    --enable-cache-report \
    --enable-metrics \
    --tokenizer-worker-num 1 \
    --host 0.0.0.0 \
    --port 30080

来源：LMSYS Blog · lmsys.org

Agent 智能体推理能力开源生态部署工程产品更新#产品更新#部署/工程#开源/仓库#推理#Agent