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

Reflection’s Beam is a 501B-parameter mixture-of-experts model with 23B active parameters per token, built for software engineering and long-horizon agentic tasks. With its planned Apache 2.0 release, Beam is designed for developers to deploy and build on across the runtimes, frameworks, and coding harnesses they already use.

SGLang now supports Beam natively and powers Reflection’s production API. It helps developers translate Beam’s inference efficiency into practical and production-proven serving performance. BF16, NVFP4, and dynamic FP8 support provide flexibility across precision and memory requirements; DFlash2 speculative decoding accelerates generation; and prefix caching with hybrid-attention memory management enables efficient reuse of context across agent interactions. Underneath these serving capabilities, we optimized Beam’s projections, normalization, and communication to improve throughput.

This post covers the integration, optimizations, and performance results behind Beam on SGLang. To deploy the model, start with our Beam serving cookbook for validated configurations and launch instructions.

Highlights

• Beam 1.0 501B-A23B: a decoder-only MoE with sliding-window and full attention, 256 routed experts and top-8 selection, and a native 1M-token context window.
• Checkpoint support: BF16, ModelOpt NVFP4, and compressed-tensors dynamic FP8.
• Serving features: DFlash2 speculative decoding, radix caching and HiCache, chunked prefill, CUDA graphs, TP/EP/PP and attention data parallelism.
• Kernel optimizations: projection packing, fusion across attention and residual boundaries, a specialized BF16x3 GEMM for FP32 router, and GB300 MoE tuning.
• Performance: On an OpenHands multi-turn agentic coding replay with the NVFP4 checkpoint on 4 NVIDIA GB300 GPUs, SGLang reaches 127.6K total tokens/sec per GPU at 185.87 tokens/sec per user (p50) with concurrency 32.

Model Architecture

Beam combines grouped-query attention with a hybrid of sliding-window and full-attention layers. Its feed-forward stack starts with dense SwiGLU layers and then switches to routed MoE. Beam gates the attention output before its projection and places sandwich normalization around both branches. Depth-scaled residuals control how each branch joins the hidden-state stream.

Beam's decoder and the operations inside one layer. Both the attention and feed-forward branches end with output normalization and a depth-scaled residual addition.

The checkpoint specifies the attention and feed-forward schedules. The native attention schedule places full attention after groups of sliding-window layers and ends with full attention. This gives the serving stack two KV lifetimes to manage: the complete history for full attention and the recent window for sliding attention.

Hybrid Attention and Output Gating

GQA shares eight KV heads among 48 query heads, reducing the KV cache that SGLang stores per token. Sliding-window layers also restrict their live history to the recent window and use RoPE. Full-attention layers use NoPE with position-dependent SSMax query scaling. Both normalize Q and K before attention.

Beam gates the attention result before the output projection. A second projection produces gate activations from the normalized hidden state; the model adds a learned per-head bias, applies sigmoid, and multiplies the attention output by the elementwise gate. The QKV and gate projections share an input, which lets us pack them into one BF16 projection in eligible layouts.

Dense Layers and Grouped MoE Routing

The first three feed-forward layers use dense SwiGLU. Later layers select eight experts from a pool of 256 and combine their SwiGLU outputs.

The router computes FP32 logits. Grouped top-k selection uses sigmoid scores and an expert correction bias to choose the experts. The backend retains FP32 routing coefficients and applies Beam's routed scaling factor during expert combination, or carries it in the routing weights when that kernel requires it.

SSMax

Beam’s full-attention layers use Scalable Softmax (SSMax) to keep attention focused as context grows. After Q/K normalization, each query at zero-based position p is scaled by (1+(p+1)/2048)0.1. This gradually sharpens the attention distribution, counteracting the tendency for softmax weights to spread across more tokens in long contexts.

Sandwich Normalization and Depth-Scaled Residuals

Each attention and feed-forward branch has input and output RMS normalization. Beam scales the branch and residual by coefficients derived from layer depth, then adds them. Both scaled operands round to the activation dtype before the addition.

An unfused transition writes the normalized branch, reads it for residual combination, and reads the combined result for the next input norm. These memory passes repeat through the decoder, so we fuse the transition while keeping Beam's rounding boundaries. The final decoder norm and language-model head feed FP32 logit softcapping.

Feature Support

Dflash2 Speculative Decoding w/ Data-Parallel Attention

Beam supports DFlash2 through SGLang's DFLASH runtime. The target captures selected decoder layers' post-layer residuals before final normalization. The drafter uses those features to propose a token block, and Beam verifies that block in one target forward.

Beam's integration adds a new CUDA serving path that combines DFlash2 with attention data parallelism. With --attn-dp-size N and --enable-dp-lm-head, the runtime keeps the attention and draft-head work aligned with the DP layout. SGLang initializes target-verification metadata after DP padding establishes the final token shape, including batches with idle peers.

We include draft KV memory in the hybrid pool budget so target-cache allocation leaves room for the drafter. Draft capture requires PP=1, and the serving context must fit both the target and the chosen drafter.

Unified Memory for Target and Draft KV

Beam's full-attention KV grows with context, while its sliding layers retain a recent window. Upstream SGLang's unified KV allocator backs these allocations with a shared GPU page budget. Request tables keep stable virtual token IDs; attention reads and KV writes translate them into the current physical slots. The translation interface uses the same virtual-to-physical contract for both directions.

DFlash admission to the unified pool extends this layout to eligible full-attention drafters. Target and draft keep distinct KV rows inside the target's page envelope, so the allocator accounts for both together. A compact-window drafter can retain its own request-to-token table over the target's virtual IDs, limiting which pages it reads while its KV storage stays in the shared layout. Draft writes translate those IDs before storing KV.

Placement depends on draft geometry, KV dtype and attention backend. The gate declines shared draft placement with HiCache, host-pool retraction or decode context parallelism. A hybrid-SWA host budgets the private draft fallback; unsupported host/draft combinations are rejected at startup.

We extend this path to trtllm_mha targets and drafters with decode CUDA graphs. It uses page tables planned for the current iteration's read/write extent and zeros new page envelopes. On SM90/SM120, it separates the fmha_v2 prefill workspace from XQA decode semaphores. The extension carries translated KV locations through target verification and draft execution in the TRTLLM backend.

Heterogeneous TP for Prefill/Decode Disaggregation

Prefill and decode process different token-batch shapes, so a disaggregated server can choose a separate attention TP size for each phase. For Beam's hybrid attention, that handoff must preserve full-attention KV and the live sliding window as KV-head ownership changes between the two groups.

PR #40909 adds this transfer path for hybrid sliding-window models through NIXL and Mooncake. It copies window state in token/head segments using the existing rank-ownership mapping, placing each token's head segment into the matching decode buffer. Staging carries full-attention KV; window state follows the head-sliced path.

Reasoning, Tool Calls and LoRA

SGLang reads Beam's checkpoint response template to separate reasoning, answers and tool calls for streaming and non-streaming requests. The integration covers Chat Completions, Responses and Messages, with Beam tool grammars and JSON-output constraints.

BF16 LoRA serving uses the Triton MoE runner and covers attention projections and gates, dense MLPs, routed experts, embeddings and the output head. The loader keeps the attention gate's qualified name so it does not collide with an MLP gate target.

BF16, NVFP4 and FP8 Model Support

We support the native BF16, FP8, and NVFP4 checkpoints. We add an NVFP4 expert-parallel alternative through FlashInfer MegaMoE. It requires both the flashinfer_megamoe runner and all-to-all backend; BF16 and dynamic FP8 remain outside this Beam extension. The optimization section explains its token ownership and communication path.

The router retains its precision and routing-scale contract across these formats. Quantized checkpoints and LoRA keep separate QKV and attention-gate modules for their loading and adapter paths.

Prefix Reuse, Chunked Prefill and HiCache

A cached Beam prefix contains full-attention history and the sliding-window state needed at its reuse boundary. SGLang aligns both with the paged cache so a request can resume from the prefix with its live window intact. Chunked prefill processes long prompts in parts alongside decode work.

HiCache extends prefix reuse from GPU memory to host memory and supported storage backends. Beam keeps its hybrid SWA pool and sliding-window memory savings with HiCache on. Host load-back uses each pipeline stage's local layer mapping.

Optimizations

Fusion Optimizations

Beam's normalization and gating paths chain small operations around the projections and collectives. We fuse these chains to remove intermediate stores, copies and launches.

• Kernel Fusion across projections, normalization, and collective boundaries. We fuse all-reduce, output norm, scaled residual addition and next input norm, extending eligible boundaries to NVFP4. PDL lets the following norm load weights before its dependency wait. Sliding layers fuse Q/K normalization with RoPE. Full-attention layers use fused Q/K normalization and reuse SSMax position scales. Eligible BF16 layers pack QKV and gate projections into one launch. Attention gating fuses bias addition, sigmoid and multiplication. The BF16 MoE collective path also incorporates local expert summation, routed scaling and BF16 rounding.

• Combine expert communication and compute with MegaMoE. We enable Beam's NVFP4 experts on FlashInfer MegaMoE, pairing its runner with the MegaMoE all-to-all backend. FlashInfer runs expert dispatch, computation and combine inside its fused EP kernel. Beam supplies explicit expert IDs and routing weights and receives complete outputs for each rank's tokens. The fused transitions preserve the rounding of both scaled residual operands before addition. Attention and MoE have separate fusion eligibility checks; MoE reduction follows the selected backend's ownership. CPU offload, empty pipeline stages and unsupported collective shapes retain the regular paths. All EP ranks participate in expert execution, including ranks with no local tokens.

MegaMoE token ownership with data parallel attention. Dense MLP layers retain TP gather/scatter; sparse layers pass DP-local tokens through fused expert dispatch, computation and combination. Complete outputs return to each token’s owning rank without an extra all-reduce.

BF16 MoE Kernel Optimization

Prefill and decode give the expert kernels different token shapes. We added GB300 tile configurations for the gate/up and down projections, with TMA for selected prefill shapes. The Triton tuner also reads Beam's expert dimensions from its configuration.

For single-token decode, an opt-in path computes routed expert outputs with two GEMV kernels. It preserves the gate/up and activation BF16 rounding boundaries and applies FP32 routing weights before rounding the down-projection output. The caller still performs expert reduction, routed scaling and TP communication. The FP32 reduction order differs from the tensor-core path.

TPU Support

We’re bringing day-0 TPU support to Beam through SGLang-TorchTPU, a native PyTorch backend for LLM, VLM, and diffusion model inference on TPUs. It combines TPU acceleration with SGLang features such as RadixAttention and HiCache to enable cost-efficient, production-ready serving for agentic workloads.

TPU recipes for Beam and access to SGLang-TorchTPU will be available soon. Sign up through this Google Form to receive the latest updates.

Reinforcement learning in Miles

Miles supports Beam through a Megatron plugin and runs rollouts on Reflection's SGLang build. The training backend implements Beam's sandwich normalization with depth-scaled residuals, SSMax attention scaling, hybrid sliding-window attention, gated attention output, auxiliary-loss-free MoE routing, and FP32 final-logit softcapping, and loads the Hugging Face checkpoint directly through Megatron-Bridge. The goal is to keep trainer and rollout log-probabilities for the same responses as close as possible.

Parallelism. The backend supports DP, TP, SP, EP, expert TP, and PP. TP shards the attention projections, including the output gate, and SP shards the activations between them along the sequence. EP and expert TP shard the routed experts independently of the attention layout, and pipeline stages can hold unequal numbers of layers. With routing replay enabled, the trainer–rollout gap stays at the same level across these layouts and does not grow over training steps.

Routing replay. Rollout Routing Replay is available as an opt-in. SGLang records the experts selected for each generated token, and the trainer's MoE layers reuse those assignments instead of recomputing top-k, so near-tied router scores cannot select a different expert path in training. Replayed routes also apply during activation recomputation and on every pipeline stage.

Numerical consistency. The trainer follows SGLang's numerics wherever Beam is precision-sensitive. The final-logit softcap runs in FP32, and depth-scaled residual additions round at the same points as the engine's fused kernel. The router gate stays in FP32 through training, checkpointing, and weight synchronization. Together these changes halve the trainer–rollout log-probability gap, leaving it at the level of BF16 reduction-order noise. Deterministic settings make repeated runs reproducible.

Colocated and disaggregated deployment. Beam runs in either mode. Disaggregated runs place training and rollout on separate GPUs and overlap the next rollout with the current training step; TIS mitigates the off-policy mismatch resulting from this one-step policy staleness. Colocated runs share the same GPUs and alternate synchronously, offloading trainer state to node-local NVMe while the engines generate. In both modes, updated BF16 weights are transferred to the engines in buckets after each step.

LoRA Training Support. LoRA training goes through Megatron-Bridge PEFT and covers attention, dense MLP, and routed-expert projections. Adapters are exported in Hugging Face format and served directly by SGLang. During RL, only adapter weights are streamed to the engines through the weight-update session while base weights stay resident, and the engines verify checksums of each received adapter.

Validated Run

The figures below show a full-parameter GRPO run on DAPO-Math-17K, starting from an early Beam checkpoint, in the disaggregated asynchronous mode on 96 H200 GPUs. 64 GPUs train with TP8, SP, EP8, expert TP8, and DP8, and 32 serve two TP16 SGLang engines. Each rollout generates 32 prompts × 8 responses with a 4K-token response cap and feeds one optimizer step. Over the first 204 steps, AIME25 accuracy rises from 46.3% to 56.7% while mean response length decreases. During the run, training resumed once from a saved checkpoint.

Performance Results

OpenHands Agentic Coding Sweep

We collected performance results for Beam NVFP4 on four GB300 GPUs with TP4 using an OpenHands multi-turn agentic coding workload. Each conversation starts with an around 74K token prompt and outputs around 220 tokens each turn, for a total of 13 turns. Successive turns reuse the prefix for a ~92% aggregate cache hit rate. We sweep concurrency across 1, 4, 8, 16, and 32 conversations and plot per-GPU token throughput (tok/s/GPU) against p50 interactivity (tok/s/user). Total throughput includes cached input tokens. These measurements use eight-token DFlash2 blocks with simulated acceptance length fixed at four and real draft tokens; they do not measure natural draft acceptance or model accuracy.

AgentX Agentic Coding Sweep

In collaboration with NVIDIA, we collected performance results for Beam NVFP4 on four GB300 GPUs with TP4/EP1 using AgentX, SemiAnalysis's replayer for recorded multi-turn agentic coding traces. AgentX reproduces the prompt sizes, prefix sharing, subagent fan-out and inter-turn timing of real coding-agent sessions, so the results reflect "realistic agentic load". Requests in a session share a long prefix, so 92% to 96% of input tokens hit the prefix cache. FP8 KV cache was utilized, and the selected speculative block size was 8. We sweep concurrency from 1 to 96 sessions and plot per-GPU token throughput (tok/s/GPU) against p90 interactivity (tok/s/user). At concurrency 96, throughput reaches 83,504 tok/s/GPU at 43.9 tok/s/user; a single session runs at 348.0 tok/s/user.

Acknowledgments

We would like to express our gratitude to the following organizations and individuals, for their contribution to the development of the SGL integration of Beam, and for extensively benchmarking its performance.

SGLang Community/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

Appendix

OpenHands Reproduction

The scripts behind the Beam OpenHands sweep are in benchmark/beam_openhands on the public bench-openhands-gb300-sweep-beam branch. Start from a prepared Beam docker image, then use reproduce.py to install the isolated client, build the dataset, and run concurrency levels 1, 4, 8, 16, and 32 on one 4 × NVIDIA GB300 node, using TP4/EP1.

The server is SGLang from that branch. The client is EvalScope at commit acd09b44, using its swe_smith multi-turn workload with the Beam replay adapter. HiCache is configured at 88 GB per rank for concurrency 16 and 32, and disabled at lower concurrency.

The workload is an OpenHands multi-turn agentic replay built from nebius/SWE-rebench-openhands-trajectories and padded with nvidia/OpenScienceReasoning-2 text: 13 turns per conversation, 74,160 bare input tokens on the first turn, 753 additional input tokens on each later turn, and 220 output tokens per turn. Growing conversation history produces approximately 80K prompt tokens per request and a 92% aggregate prefix-cache hit rate.

The server launch commands are as follows, with <cc> setting both the maximum running requests and the decode CUDA-graph batch cap

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