Repository navigation
Conversation
Add a Beam serving image on top of the v0.5.22 release image. It clones the upstream beam branch and adds the pins that branch needs, without changing docker/Dockerfile. Split from #44747.
…ary (2/5) Add the Beam fused projection/norm/residual and attention-gate ops, the FP32 router GEMV, BF16x3 and cuTile router GEMMs, the Beam all-reduce boundary kernel, and their kernel tests and benchmarks. Split from #44747. Co-authored-by: BBuf <1182563586@qq.com> Co-authored-by: Caihua Li <caihua.li@bytedance.com> Co-authored-by: Florent Gbelidji <florentgbelidji@Florents-MacBook-Pro.local> Co-authored-by: iaroslav-reflection <iaroslav@reflection.ai> Co-authored-by: Jinze Xue <jinzex@nvidia.com> Co-authored-by: mloh <mloh@nvidia.com> Co-authored-by: neilkg <33635204+neilkg@users.noreply.github.com> Co-authored-by: thibautbar <t.barroyer@icloud.com>
Beam serving does not use Dynamo, and the v0.5.22 release already resolves grpcio-tools 1.84.0. The pins move to their own PR so they can be dropped without touching this layer.
The kernel, its test and its benchmark need the tileiras compiler. They move with the router dispatch that uses them, so this layer needs no new dependency.
The image is #45253 on its own; nothing in this PR uses it.
main has the merged #42930 version.
…shared-memory write
github-actions
Bot
added
documentation
Jiminator
force-pushed
the
beam-stack-02-kernels
branch
from
October 25, 2026 09:48
98141d2 to
851658f
Compare
Jiminator
force-pushed
the
beam-sm103-moe-gemv
branch
from
October 25, 2026 09:48
5808f6a to
15944d9
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Move the opt-in SM103 single-token BF16 MoE GEMV out of #45254 so the Day-0 kernels can land independently. This draft keeps the GEMV, its shared SiLU helper and focused checks together. NVFP4 end-to-end validation does not exercise this BF16 path.
The operator uses Beam's TP4 expert shape (
256experts, hidden size5120, intermediate size640, top-k8) and returns weighted per-expert BF16 outputs. The caller owns expert reduction, routed scaling and TP communication. Default model dispatch is unchanged here.Validation:
torch 2.14.1+cu130, Triton3.8.0):1 passed, 2 warnings in 1.24s. It checks an independent FP64 reference with BF16 rounding points, int32/int64 expert IDs, and five graph replays with changed inputs. The kernel, helper and test bytes match the tested source.See the run instructions and focused correctness test.
Base: #45254 (
beam-stack-02-kernels); retarget tomainafter that parent lands. The existing opt-in dispatch flag in #45256 needs coordination when the model stack is restacked. This draft does not depend on #45361.CI States
Latest PR Test (Base): ❌ Run #40182158707
Latest PR Test (Extra): ❌ Run #40182158602
Latest PR Test (AMD ROCm 10): 🚫 Run #40182158433