Repository navigation
[Feat] beam: Add Beam model, runtime support and checkpoints (4/7) - #45256
Conversation
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.
| } | ||
| 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( |
There was a problem hiding this comment.
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)
There was a problem hiding this comment.
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"): |
There was a problem hiding this comment.
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
| full_cell_size = ( | ||
| self._full_cell_bytes() | ||
| + self._swa_per_token * self._draft_swa_full_layers_num | ||
| + self._draft_cell_size |
There was a problem hiding this comment.
Is it correct to add draft size here?
There was a problem hiding this comment.
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( |
There was a problem hiding this comment.
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. | ||
| """ | ||
|
|
There was a problem hiding this comment.
Can you split the part of megamoe autotune optimziation into a standalone PR, since it's not specific to Beam
| ) | ||
| if ( | ||
| head_dims in cutlass_supported_head_dims | ||
| and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE) |
There was a problem hiding this comment.
Can we revert the changes in this PR? Since beam only applies trtllm mha backend in cookbook.
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.
…(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>
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:
Beam model and config:
DFlash2BeamDraftModelalias, 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.E=256,N=640, used by the BF16 Triton MoE path.SGLANG_OPT_BEAM_DECODE_GEMV).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.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_BACKENDare 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-metricsfails during startup warmup withTypeError: '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
NotImplementedErrorfor--enable-loraand--lora-pathswhen the model is built. LoRA serving will follow separately.The split commits carry the
Co-authored-bytrailers 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.
reflection/Beam-1.0-501B-A23B-NVFP4,reflection/Beam-1.0-501B-A23B-FP8andreflection/Beam-1.0-501B-A23Bcheckpoints, plus the DFlash2 draft.dfa85fe771b3bcf014f5af5913f2ed363d31dfdfpasses all 133 required e2e cases at temperature 0, with its source and fixture integrity check passing and no skips./tokenizeall pass.silu(gate) * up, and at the routing-weighted output.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
mainandmainis merged into it. #44747 stays open, unchanged, as the full reference. Merge in this order:nvidia-cuda-tileiras, needed by [Perf] beam: Add optional cuTile BF16x3 router GEMM #45361#45292 (Dynamo and
grpcio-toolspins 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