Skip to content

[Feat] beam: Add Beam model, runtime support and checkpoints (4/7) - #45256

Merged
Jiminator merged 66 commits into
beam-stack-02-kernelsfrom
beam-stack-04-model
Oct 25, 2026
Merged

Jiminator merged 66 commits into
beam-stack-02-kernelsfrom
beam-stack-04-model

Conversation

@Jiminator

@Jiminator Jiminator commented Oct 19, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

Split #44747 into reviewable parts. This part adds the runtime support Beam relies on and the Beam model, position 4 of 7 in the merge order below. It folds in #45255, which is closed in its favor, the way Muse Spark's runtime and model landed together. It is based on #45254's branch (beam-stack-02-kernels), so its diff shows only this part.

Modifications

Runtime support:

  • TRT-LLM MHA. An extend batch whose trailing requests extend by one token runs that tail on the decode kernel. This now also covers FP8 KV and batches made only of one-token requests. It extends [Perf] Route BF16 mixed-batch decode tails through decode attention #43238 and runs on trtllm-gen (SM100) only; SM90 and SM120 take FMHA-v2 or the context-parallel path and never reach it.
  • FlashInfer MegaMoE. Beam is added to the validated MegaMoE architectures. The startup autotune fix that the high-throughput recipe needs with default autotune, and the optional decode workspace, are in [Fix] flashinfer-megamoe: Fix startup autotune on FlashInfer 0.7.2 and add an optional decode workspace #46213.
  • PD decode. A token-ID logprob row is reserved for the handoff token, so streamed rows stay aligned. That row is empty, because prefill does not send token-ID logprobs.
  • Forward-pass metrics. DFlash is added to the spec decode pre-settlement snapshot.
  • Pool configurator. The SWA request-cap budget now includes the DFlash draft cell.

Beam model and config:

  • Hybrid sliding-window/full attention, SSMax query scaling, attention output gating, sandwich norms, depth-scaled residuals and grouped MoE routing.
  • BF16, ModelOpt NVFP4 and compressed-tensors FP8 checkpoints. Compressed-tensors NVFP4 is rejected.
  • DFlash integration. Beam uses the general CUDA attention-DP implementation from merged [Feat] spec: Support CUDA attention DP for DFlash and DFlash2 #46330 and draft placement from merged [Refactor] spec: Resolve draft parallel placement from model declarations #46290. This PR adds the DFlash2BeamDraftModel alias, Beam's embedding/lm-head DP mapping, and hidden-state capture. The generic DFlash admission hook, worker, embedding helper and verify-metadata path match the parent branch.
  • GB300 Triton MoE configs for E=256,N=640, used by the BF16 Triton MoE path.
  • Opt-in single-token MoE GEMV (SGLANG_OPT_BEAM_DECODE_GEMV).
    • It uses the unquantized Triton MoE path and requires SGLANG_OPT_FUSE_SWIGLU_INTERLEAVED=1, BF16 experts of the Beam shape, MoE TP4/EP1 and SM103. Each invocation also needs one contiguous BF16 token.
    • Other setups, and forwards with more than one token, keep their normal MoE path. Only a failed static check logs that the switch is unused.
    • It is not bit-exact with the Triton path.

Beam's router uses the tiny FP32 GEMM up to 16 tokens for BF16 checkpoints, then the GEMV, then the cuBLAS BF16x3 GEMM. The cuTile router and SGLANG_BEAM_ROUTER_BACKEND are in #45361.

General fixes found while adding Beam are in their own PRs:

The only test change is DFlash in test_forward_pass_metrics.py's snapshot loop. Without the metrics change, Beam with DFlash and --enable-forward-pass-metrics fails during startup warmup with TypeError: 'NoneType' object is not iterable.

The Beam E2E test is in #45257. The unchanged upstream DFlash admission, embedding and metadata tests come from #46330.

LoRA serving for Beam is not part of this split. Beam raises NotImplementedError for --enable-lora and --lora-paths when the model is built. LoRA serving will follow separately.

The split commits carry the Co-authored-by trailers of the Beam commits that touched these files.

Validation

Current DFlash integration cleanup (903d45e17d): 58 tests passed across seven existing DFlash/draft-placement files, with one existing CUDA-only test skipped in the CPU-only run. The nine admission/embedding/metadata cases pass unchanged. No test code was added or modified. The two remaining generic DFlash overrides were removed so this branch uses #46330's implementation.

Historical GPU/model results: the runs below predate this cleanup and were not repeated for it.

All runs are on 4×GB300 with FlashInfer 0.7.2. The tree is this PR plus #43216, which has since merged.

Configuration GSM8K Notes
Cookbook low-latency NVFP4 (TP4, DFlash, without parser flags) 96.36% zero request errors; the DFlash2 selector is folded into the draft CUDA graph
Cookbook low-latency FP8 96.66% zero errors
Cookbook high-throughput NVFP4 (MegaMoE TP4/DP4/EP4), cold autotune cache 96.44% with #46213's autotune fix; MegaMoE autotune completes
High-throughput NVFP4 plus DFlash, default autotune 96.44% with #46213's autotune fix; zero errors
The same with autotune off 96.44% zero errors
BF16 TP4, attention DP2, DFlash 96.44% concurrent outputs identical to isolated in the probe
BF16 TP4 with the decode GEMV 96.21% the GEMV kernels run at batch size 1
Cold FlashInfer JIT (no AOT top-k), DFlash with attention DP4 96.51% the selector's top-k builds during startup; without the warm-up, the first decode builds it inside the scheduler watchdog and the server is killed
  • Verifier. With [Feat] parser: Add Beam reasoning, tool calls and response templates (5/7) #45257's parser changes on top of this branch, the pinned Beam verifier dfa85fe771b3bcf014f5af5913f2ed363d31dfdf passes all 133 required e2e cases at temperature 0, with its source and fixture integrity check passing and no skips.
    • The run uses the cookbook low-latency recipe with parser flags and scores 96.59% on GSM8K with zero request errors.
    • Chat reasoning, streaming, a chat tool call, Responses text, a Responses function call, a Responses shell round trip and /tokenize all pass.
    • No optional verifier suite or third-party evaluation was run.
  • TRT-LLM mixed tail. The mixed-tail cases from [Feat] beam: Add native Beam serving support #44747 pass (16 passed, 2 skipped on SM100). They include batches of only one-token requests, split-versus-unsplit parity, long-context FP64 references and CUDA-graph replay.
  • GEMV numerics. Both kernels were compared with an FP64 reference that rounds to BF16 where the kernels do: after the gate/up projection, after silu(gate) * up, and at the routing-weighted output.
    • On 8 real checkpoint experts, the GEMV's BF16 output matches the reference on nearly every element, with relative L2 below 2e-9.
    • In the tested random-input cases, the GEMV is within relative L2 3.4e-4.
    • The Triton deferred-sum path is within 4.2e-4 to 6.8e-4 of the same reference, and the GEMV versus Triton is within 6.3e-4.
    • With a fixed kernel config, the shared Triton MoE path is bitwise unchanged from the base.
  • Existing GPU tests. Main's tests for the touched code pass unedited: TRT-LLM MHA dense, graph metadata, encoder-only and stride parity; Triton MoE padding and fused MoE; and the DFlash buffer, Domino and graph-pool tests.
  • CPU tests. The CPU tests that import the touched modules pass, except three files that fail the same way on main in this environment. pre-commit passes.
  • Known behavior. Under attention DP4, a request's greedy output can change with other ranks' load.
    • This was observed with MegaMoE, with the TRT-LLM MoE runner without MegaMoE, without DFlash, and on the previously published head.
    • In the probe's valid comparisons, the isolated run's top-2 logprob margin at the first differing token ranged from 0 to 0.262 nat.
    • Isolated requests were deterministic in repeat probes, and these runs showed no GSM8K accuracy regression and zero request errors.
    • This does not establish universal batch invariance or a specific kernel-level cause. BF16 with attention DP2 matched isolated outputs exactly.
  • Not validated here:
    • NPU;
    • Hopper and SM120 (all runs are on GB300);
    • BF16 on two nodes (TP8), including the TP8 cluster path of [Feat] kernels: Beam model support #45254's all-reduce boundary;
    • serving with PD token-ID logprobs (that fix has CPU checks only).

Merge order

Each part is based on the one before it, so its diff shows only its own changes. Merge bottom-up; after each part merges, the next part is retargeted to main and main is merged into it. #44747 stays open, unchanged, as the full reference. Merge in this order:

  1. [Chore] docker: Add Beam CUDA image (1/8) #45253 — Beam CUDA Docker image (merged)
  2. [Chore] deps: Bump the OpenAI SDK to 2.8.0 for Responses shell tools (2/8) #45291 — OpenAI SDK 2.8.0, needed by the Responses shell tools in [Feat] parser: Add Beam reasoning, tool calls and response templates (5/7) #45257 (merged)
  3. [Feat] kernels: Beam model support #45254 — Kernels
  4. [Feat] beam: Add Beam model, runtime support and checkpoints (4/7) #45256 — Beam model and runtime integration (using upstream [Feat] spec: Support CUDA attention DP for DFlash and DFlash2 #46330 for DFlash CUDA DP) ([Feat] srt: Add attention, memory and DFlash DP support for Beam (4/8) #45255 folded in; also after [Fix] flashinfer-megamoe: Fix startup autotune on FlashInfer 0.7.2 and add an optional decode workspace #46213 for the high-throughput recipe with default autotune; [Fix] nccl: Disable graph buffer registration whenever a captured graph can hold a collective #43216 is merged)
  5. [Feat] parser: Add Beam reasoning, tool calls and response templates (5/7) #45257 — Beam reasoning, tool calls, response templates and the Beam E2E test ([Feat] parser: Enrich chat_parsing streaming events #40478, [Fix] parser: Accept whole-number floats for integer tool arguments #43941 and [Feat] parser: Add checkpoint-driven response-template adapters #40479 are merged)
  6. [Chore] deps: Add nvidia-cuda-tileiras for the optional cuTile router #45290 — nvidia-cuda-tileiras, needed by [Perf] beam: Add optional cuTile BF16x3 router GEMM #45361
  7. [Perf] beam: Add optional cuTile BF16x3 router GEMM #45361 — cuTile BF16x3 router GEMM (optional: without it Beam uses the cuBLAS router)

#45292 (Dynamo and grpcio-tools pins for the Beam image) is not in the merge order for now.


CI States

Latest PR Test (Base): 🚫 Run #40237989506
Latest PR Test (Extra): 🚫 Run #40237989484
Latest PR Test (AMD ROCm 10): 🚫 Run #40237989472

Jiminator and others added 2 commits October 19, 2026 16:50
Add the shared runtime support Beam relies on: TRT-LLM MHA and FlashInfer attention fixes, unified hybrid SWA memory, DFlash with attention DP, FlashInfer MegaMoE all-to-all under attention DP, and their tests.

Split from #44747.

Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Florent Gbelidji <florent@reflection.ai>
Co-authored-by: Florent Gbelidji <florentgbelidji@Florents-MacBook-Pro.local>
Co-authored-by: Florent Gbelidji <florentgbelidji@florents-macbook-pro1.home>
Co-authored-by: Florent Gbelidji <florentgbelidji@Host-002.lan>
Co-authored-by: iaroslav-reflection <iaroslav@reflection.ai>
Co-authored-by: Jiajun Li <jiajun.li@radixark.ai>
Co-authored-by: Jinze Xue <jinzex@nvidia.com>
Co-authored-by: Kaiwu <kai@reflection.ai>
Co-authored-by: mloh <mloh@nvidia.com>
Co-authored-by: Neil Gollapudi <neil@reflection.ai>
Co-authored-by: neilkg <33635204+neilkg@users.noreply.github.com>
Co-authored-by: Qiaolin-Yu <liin1211@outlook.com>
Co-authored-by: Renault Young <renault@reflection.ai>
Co-authored-by: thibautbar <t.barroyer@icloud.com>
Co-authored-by: Yoni Gozlan <yoni@reflection.ai>
Co-authored-by: zhihengy <andy.ye@radixark.ai>
Add the Beam model and config, BF16/ModelOpt NVFP4/compressed-tensors FP8 checkpoint support, the DFlash2 Beam draft alias, GB300 MoE tuning, and the registered unit and end-to-end tests. Beam rejects LoRA adapters here.

Split from #44747.

Co-authored-by: AymenKallala <aymen.kallala@gmail.com>
Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Florent GBELIDJI <florent@reflection.ai>
Co-authored-by: Florent Gbelidji <florent@reflection.ai>
Co-authored-by: Florent Gbelidji <florentgbelidji@Florents-MacBook-Pro.local>
Co-authored-by: Florent Gbelidji <florentgbelidji@florents-macbook-pro1.home>
Co-authored-by: Florent Gbelidji <florentgbelidji@Host-002.lan>
Co-authored-by: Florent Gbelidji <florentgbelidji@Mac.lan>
Co-authored-by: iaroslav-reflection <iaroslav@reflection.ai>
Co-authored-by: Jinze Xue <155670984+jinzex@users.noreply.github.com>
Co-authored-by: Jinze Xue <jinzex@nvidia.com>
Co-authored-by: kai <kai@reflection.ai>
Co-authored-by: mloh <mloh@nvidia.com>
Co-authored-by: mreso <13337103+mreso@users.noreply.github.com>
Co-authored-by: Qiaolin-Yu <liin1211@outlook.com>
Co-authored-by: Renault Young <renault@reflection.ai>
Co-authored-by: thibautbar <t.barroyer@icloud.com>
Co-authored-by: Yanbin Jiang <jybsuper@gmail.com>
Co-authored-by: Yoni Gozlan <yoni@reflection.ai>
Restore the router dispatch from before the cuTile BF16x3 GEMM: the
tiny FP32 GEMM up to 16 tokens for BF16 checkpoints, then the GEMV,
then the cuBLAS BF16x3 GEMM. SGLANG_BEAM_ROUTER_BACKEND and the
router dispatch tests move with the cuTile router.
@Jiminator Jiminator changed the title [Feat] beam: Add Beam model, checkpoints and tests (4/5) [Feat] beam: Add Beam model, checkpoints and tests (5/8) Oct 20, 2026
}
fake_modules["sglang.srt.environ"].envs = types.SimpleNamespace(
SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=types.SimpleNamespace(get=lambda: False),
SGLANG_FLASHINFER_MEGAMOE_MAX_TOKENS_PER_RANK=types.SimpleNamespace(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the extra unit tests, please remove all the unnecessary ones.
We need to add an e2e test for Beam (maybe at extra groups or nightly groups)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the extra unit tests.

The one left adds DFlash to the snapshot loop in test_forward_pass_metrics.py. It guards a real crash where Beam’s low-latency recipe with --enable-forward-pass-metrics dies in startup warmup because trtllm_mha doesn’t keep CPU seq lens on spec decode batches.

The Beam e2e test is test_beam_10_b200_pd.py in #45257

)

if attn_dp_enabled_of(cfg) and cfg.device != "npu":
if not cfg.device.startswith("cuda"):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems dp attention + dflash is open to every model. Can we add a guard that throws error when detecting the model is not Beam here

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

full_cell_size = (
self._full_cell_bytes()
+ self._swa_per_token * self._draft_swa_full_layers_num
+ self._draft_cell_size

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it correct to add draft size here?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should be correct, the DFlash draft pool reuses the target’s token count, so each full-pool token also costs _draft_cell_size. HybridSWAPoolConfigurator already counts it; the request-cap override dropped it, so the full pool takes the whole budget and the draft pool overflows it.

self.layer_num = self.full_layer_nums + self.swa_layer_nums
self.start_layer = 0
# PP stages keep global layer ids; HiCache load events count from the first.
self.start_layer = min(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do this change?

@Jiminator Jiminator Oct 24, 2026 •

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed, previously on an old version of the fork, PP=2 + HIcache was broken without it, but its no longer needed. There is a very related fix opened in #46200 for a seperate bug with PP

communication internally via the deep_gemm symmetric buffer, so the dispatcher
and combine stay pure no-ops; this module owns the layer build + forward.
"""

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you split the part of megamoe autotune optimziation into a standalone PR, since it's not specific to Beam

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done and split into #46213

)
if (
head_dims in cutlass_supported_head_dims
and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we revert the changes in this PR? Since beam only applies trtllm mha backend in cookbook.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done and split into #46214

Beam selects trtllm_mha when no attention backend is set. The fix applies to every hybrid-SWA model on --attention-backend flashinfer, so it moves to a standalone PR.
Neither is specific to Beam: the startup autotune output buffer fixes every NVFP4 MegaMoE model on FlashInfer 0.7.2, and the decode workspace is an opt-in profile for any ModelOpt NVFP4 MoE. Beam stays in the validated MegaMoE architectures.
Only Beam is validated on this path; other CUDA targets are rejected at startup as before. The mRoPE check it subsumes goes away.
Jiminator and others added 4 commits October 23, 2026 19:37
…(5/7) (#45257)

Signed-off-by: Yoni Gozlan <yonigozlan@users.noreply.github.com>
Co-authored-by: Yoni Gozlan <yonigozlan@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com>
Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Jinze Xue <jinzex@nvidia.com>
Co-authored-by: mloh <mloh@nvidia.com>
Co-authored-by: Neil <neil@reflection.ai>
Co-authored-by: neilkg <33635204+neilkg@users.noreply.github.com>
Co-authored-by: thibautbar <t.barroyer@icloud.com>
Co-authored-by: Yoni Gozlan <yoni@reflection.ai>
@Jiminator Jiminator added run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) labels Oct 25, 2026
@Jiminator
Jiminator merged commit e58899d into beam-stack-02-kernels Oct 25, 2026
34 of 87 checks passed
@Jiminator
Jiminator deleted the beam-stack-04-model branch October 25, 2026 22:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

blackwell SM100/SM120 bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) dependencies Pull requests that update a dependency file documentation Improvements or additions to documentation hicache Hierarchical Caching for SGLang jit-kernel memory-pool quant LLM Quantization run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) speculative-decoding unified-radix-cache

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants