Reflection Beam 1.0 on vLLM: Day 0 Serving, Agent-Optimized Kernels, and DFlash2 Speculative Decoding

vLLM added day 0 serving support for Reflection's Beam, a 501B-parameter mixture-of-experts model with 23B active parameters per token, and optimized it specifically for NVIDIA Blackwell hardware. The result matters because it combines full production serving features (reasoning, tool calling, structured outputs, tensor/data/expert parallelism) with a kernel stack that coding agents optimized autonomously, making batch-1 decode 2.71x faster and prefill 1.87x faster than vLLM's first working implementation.

A single request reaches 708 tokens per second on TP8 using Beam's DFlash2 draft head, with an acceptance length of 3.7. Adaptive verification adds 10% more throughput at high concurrency.

The model can be served with the Docker image 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 model architecture: 52 layers, hybrid attention, DeepSeek-V3-style MoE

Beam is a 52-layer MoE decoder with a hidden size of 5120. Its MoE routing follows DeepSeek-V3; the architecture's distinctive choices are in the attention layers and in how the residual stream is normalized.

Attention: sliding-window layers plus NoPE full-attention layers

Beam splits attention into two regimes. 43 layers use sliding-window attention over the last 512 tokens with RoPE. The remaining 9 layers — every sixth layer, plus the last — attend over the full context with no positional encoding (NoPE).

Because attention flattens as the number of keys grows when RoPE is absent, Beam applies SSMax on the NoPE layers. SSMax scales each query by (1 + p/2048)^0.1, where p is the 1-based position, so the softmax sharpens as the context grows.

Beam uses grouped-query attention with 48 query heads, 8 KV heads, and head dimension 128. Q and K are RMS-normalized per head before RoPE or SSMax. Before the output projection (o_proj), each head's output is multiplied by sigma(gate + bias), where the gate comes from its own full-width projection with a learned per-head bias, so a head can be turned down or off token by token.

The attention design keeps the KV cache small. Only the 9 full-attention layers keep KV for the whole sequence: 9 * 2 * 8 * 128 values per token, which is 36 KiB in BF16 and 18 KiB in FP8. A full 1M-token context costs about 18 GiB of FP8 KV cache before sharding.

MoE and residual stream: FP32 router, sandwich norm, depth-scaled residuals

Beam has 256 routed experts with d_ff of 2,560 in NVFP4. Routing matches DeepSeek-V3: sigmoid scores plus a correction bias pick the best 4 of 8 expert groups, then the top 8 experts, and the renormalized weights are scaled by 2.5. Unlike DeepSeek-V3, Beam has no shared expert. Layers 0 through 2 use a dense SwiGLU MLP with d_ff of 20,480 instead of MoE.

The router is fully FP32: the 5120 x 256 router weights, the accumulation, and the logits all run in FP32.

Two normalization choices shape the residual stream. Every sublayer has an RMSNorm before it and another after it ("sandwich norm"). The n-th residual add computes h <- sqrt((n-1)/n) * h + x/sqrt(n) rather than h + x, so with unit-variance sublayer outputs the residual stream keeps unit variance at every depth instead of growing with it.

Agent-driven optimization made batch-1 decode 2.71x faster

vLLM's first Beam implementation was unoptimized and almost pure PyTorch. Coding agents then optimized it autonomously: they profiled the model, wrote the kernels, and benchmarked each change. On TP8, seven changes brought batch-1 decode from 8.95 ms to 3.30 ms per step, and an 8 x 2K prefill from 328 ms to 176 ms.

The seven changes were:

  1. Fusing norm, residual, and norm (-40% decode). The post-norm, scaled residual add, and next pre-norm at each of the 104 sublayer boundaries run as one Triton kernel, with FP32 math and the reference implementation's BF16 rounding points.
  2. Fusing Q/K norm, RoPE/SSMax, and the output gate. Per-head Q/K RMSNorm and RoPE (or SSMax) run as one kernel, and the sigmoid gate as another.
  3. Fusing both tensor-parallel all-reduces into the norm kernels, built on FlashInfer's multicast all-reduce. After the MoE, the same kernel also combines the top-8 expert outputs.
  4. A specialized GEMM kernel for the FP32 router. The router previously fell back to a generic path; vLLM's small-batch FP32 router kernel now supports Beam's 5120 x 256 shape while keeping FP32 weights, accumulation, and logits.
  5. Merging the QKV and gate projections with a low-latency GEMM. The Q, K, V, and gate weights are packed into one matrix at load time so the four projections run as one GEMM; for small batches it uses a low-latency CuTe DSL GEMM or FlashInfer's TinyGEMM.
  6. A fused KV-cache update. The Q/K norm and RoPE kernel also writes K and V into the KV cache, eliminating a separate cache-write kernel.
  7. Dense-MLP activation in one kernel for the three dense layers.

The latency measurements were taken as TP8 model-forward latency on GB200, for batch-1 decode and an 8 x 2,048-token prefill at roughly 100K context.

The fused decode graph runs one MoE layer in 13 kernels

After these changes, one MoE layer runs as 13 kernels in batch-1 decode on TP8. Six of those kernels were written, fused, or extended for Beam; every small Beam-specific op is merged into a neighbor, and both all-reduces live inside norm kernels.

The other seven kernels come from libraries: TRTLLM-gen attention, the cuBLAS output projection, and the TRTLLM NVFP4 MoE path (quantization, two routing kernels, and two grouped GEMMs). Most of the remaining time sits in those library kernels, which vLLM identifies as the next optimization target.

Attention DP plus MegaMoE reaches 168.8K tokens per second

In the high-concurrency region, Beam performance improves further by combining attention data parallelism with MegaMoE, powered by FlashInfer's CuTe DSL NVFP4 MegaMoE kernel, which can fuse communication with expert computation. The Finalize + norm boundary fusion remains applicable in this configuration.

Configuration Peak input token throughput
Attention DP + MegaMoE 168.8K tokens / second
TP (flashinfer_trtllm, flashinfer_nvlink_one_sided) 98.7K tokens / second

Attention DP with MegaMoE delivered 1.71x higher peak prefill throughput than TP4. The benchmark used 4 x GB300 with an input sequence length of 32K tokens and an output sequence length of 1 token to isolate prefill performance.

To enable it, replace the MoE and parallel settings with:

--moe-backend flashinfer_moe_ep_cutedsl \
  --data-parallel-size 4 \
  --enable-expert-parallel

DFlash2 speculative decoding: 708 tokens/s at concurrency 1

Beam ships with a DFlash2 draft head, which vLLM supports from day 0. Instead of drafting one token at a time, DFlash2 predicts a whole block of tokens in a single forward pass from the target model's hidden states.

The drafter is built from 5 sliding-window layers that read hidden states from 5 target layers and share the target's embedding and LM head. A grouped convolution and a small candidate selector restore some of the left-to-right dependency that a parallel drafter loses.

Adaptive verification, enabled with "enable_adaptive_verification": true, predicts how likely each draft token is to be accepted and verifies only the ones worth their cost, rather than a fixed number per step. DFlash2 has no confidence head, so the acceptance estimate comes from the drafter's own logits. The drafter's convolution, norms, and KV writes are fused, and the whole draft step runs in a CUDA graph. FP8 KV cache works for the drafter too, with no change in acceptance.

Serving results on TP8 come from the low-entropy category of SPEED-Bench's throughput_1k split: roughly 1K-token prompts, mostly code completion and text sorting, with greedy decoding and up to 4K output tokens. Low-entropy text is easier to predict, so lower acceptance should be expected on open-ended chat.

Concurrency Fixed draft (tok/s) Adaptive (tok/s) Change Acceptance length
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

Adaptive verification pays off under load. At concurrency 128 it raises output throughput by 10.8% and cuts p99 TPOT by 9%, while verifying fewer draft tokens per step; acceptance length drops slightly but the compute saved outweighs it. At low concurrency, where extra draft tokens are nearly free, the effect is neutral.

Deployment guidance: parallel setup, reasoning, and tool calling

The command at the top of vLLM's post is the recommended default: TP4 on a single GB300 node, with DFlash2 and an FP8 KV cache. To enable adaptive verification, add "enable_adaptive_verification": true to the --speculative-config. The other optimizations described above are on by default.

For the lowest latency, use TP4 or TP8. For higher throughput, use data-parallel attention with expert parallelism, for example -dp 4 --enable-expert-parallel.

The beam reasoning and tool-call parsers read the response template that ships with the checkpoint. By default, vLLM also constrains every response to that template, so tool calls always parse and match each tool's JSON schema. response_format and thinking_token_budget work as usual. To turn the default constraint off, set VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT=0.

Day 0 reinforcement learning support with SkyRL

vLLM collaborated with the SkyRL team to enable day 0 support for reinforcement learning on Beam. LoRA, BF16 full fine-tuning, and MXFP8 RL are supported out of the box, with an upgraded RL stack for the 500B scale.

The SkyRL team post-trained Beam on the DAPO recipe with vLLM for inference and Megatron for training. LoRA, full fine-tuning, and MXFP8 RL achieve similar gains in performance over 50 training steps, or 100 optimizer updates (reasoning minimal), with average accuracy (avg@8) jumping from about 61% to about 79% on AIME25. Average absolute log-probability difference between Megatron and vLLM stays under 0.01 with Rollout Routing Replay (R3) enabled for LoRA and BF16 full fine-tuning, and under 0.03 for MXFP8 RL.

Who built the Beam serving stack

vLLM describes the work as a close collaboration between the Inferact, Reflection, and Nvidia teams. Inferact contributors were Woosuk Kwon, Ziqi Zhao, Ao Shen, and Tiezhen Wang. Reflection contributors were Kaixing Wu, Yoni Gozlan, Neil Gollapudi, Thibaut Barroyer, Aymen Kallala, Matthias Reso, and Florent Gbelidji. Nvidia's contributor was Siyuan Fu.

Sources