Repository navigation
[None][feat] Extend Beam ModelingV2 inference and serving - #20952
Draft
Mgluhovskoi wants to merge 14 commits into
Draft
Mgluhovskoi wants to merge 14 commits into
Mgluhovskoi wants to merge 14 commits into
Conversation
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Add BeamDenseBlockBf16Sm103Tp1, a directly constructed component that runs
Beam decoder block 0 (sliding-window attention, dense SwiGLU MLP) in BF16 on
one sm_103 GPU, for one fresh context of 1 to 129 tokens. It is a component,
not a routed target: there is no routing.py, router-index entry or
ForCausalLM, and nothing in the engine builds it.
The forward is assembled from catalog entries (CublasMm, FlashinferRmsnorm,
FlashinferSiluAndMul, and ThopAttention over a KVCacheManagerV2 pool) and
three entries the component owns: beam_neox_rope, beam_output_gate and
beam_depth_residual. Configurations and inputs the component does not compute
are rejected before any kernel runs.
The catalog entries gain cells for the shapes the block runs. The block's
tests compare the whole block, and every operation of it on its own input,
against a plain-torch restatement of the source block at seeds {17, 41} x
lengths {1, 17, 129}; they also check the op and kernel path and that a wrong
residual coefficient is caught. They are local/QA tests and are not added to
a CI test list.
To run the tests:
env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE \
python -m pytest tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Add BeamMoeRouterBf16Sm103Tp1, a directly constructed component that runs
the sigmoid grouped top-k router of Beam decoder block 3 on one sm_103 GPU.
It takes already-normalized BF16 hidden states [T, 5120] and returns FP32
mixing weights [T, 8] and int64 expert ids [T, 8], in that order. It holds
only the FP32 router weight [256, 5120] and the FP32 correction bias [256].
It does not run experts or build a MoE block, and nothing in the engine
builds it.
The arithmetic is one owned op, BeamGroupedTopkRouter, which follows the
source router step by step in FP32 with TF32 and autocast off: projection,
sigmoid, correction bias added for selection only, group score as the sum of
the top two corrected scores in each of 8 groups of 32, 4 groups kept and
the rest masked, top-8 experts, and the unbiased scores of the chosen experts
divided by their sum plus 1e-20. The routed scaling factor is not applied.
Tied scores keep the order torch.topk(sorted=True) returns; no index
tie-break is added. The native noaux_tc op is not used because its sigmoid
form, tie rule and scaling differ from the source.
The tests run the component and compare every step of the op's call with
the op's plain-torch reference at seeds {17, 41} x tokens {1, 17, 129}, and
on controlled fixtures for the correction bias, the group and expert
selection boundaries, group exclusion, ties and the epsilon. Group and
expert ids are compared exactly and in order. They are local/QA tests and
are not added to a CI test list.
To run the tests:
env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE \
python -m pytest tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Add BeamMoeExpertBf16Sm103Tp1, a directly constructed component that runs
one routed SwiGLU expert of Beam decoder block 3 on one sm_103 GPU. It
takes already-normalized BF16 hidden states [T, 5120], T >= 1, and returns
the expert's unweighted BF16 output [T, 5120]. It holds only the three
bias-free BF16 matrices gate_proj and up_proj [2560, 5120] and down_proj
[5120, 2560], loaded by checkpoint name and copied as given. Every tensor
is checked and staged on the device before any parameter changes, so a
rejected load leaves the expert as it was. It routes no token and applies
no mixing weight, routed scaling factor, residual or collective; nothing
in the engine builds it.
The three projections are the catalog's cublas_mm, each reading the
stored weight's transpose as its column-major operand. The activation is
one owned op, SwigluFp32: silu(gate.float()) * up.float(), rounded once to
BF16 before the down projection, which also returns the FP32 SiLU and
product so each step can be observed. The shared flashinfer_silu_and_mul
entry is not used because it returns only the rounded activation.
The tests run the component at seeds {17, 41} x tokens {1, 17, 129} and
compare every step of the same forward call (gate, up, FP32 SiLU, FP32
product, BF16 activation, output) with the source expert written in plain
torch, elementwise within one BF16 unit relative plus a per-step absolute
allowance. Exact fixtures cover zero input, negative, zero and positive
gates, a precision fixture on which every activation differs from one
computed with a BF16 SiLU and product, asymmetric one-term feature maps
that expose weight orientation and the gate/up/down mapping, token-row
permutation, and zero down weights. They also check construction, exact
loading, rejected inputs and loads, and that the forward runs on the GPU.
The op's test drives its cells against an FP64 reference. They are
local/QA tests and are not added to a CI test list.
To run the tests:
env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE \
python -m pytest tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Add beam_dispatch_combine, a directly called component that dispatches
the tokens of Beam decoder block 3 to their eight routed experts and
combines the weighted expert outputs on one sm_103 GPU. It takes BF16
hidden states [T, 5120], T >= 1, the router's int64 expert ids and FP32
mixing weights [T, 8], and an expert provider, and returns BF16
[T, 5120]. It holds no state: no router parameter, no expert matrix and
no configuration object. Nothing in the engine calls it.
The provider owns the expert computation: provider(expert_id, rows) is
called once per selected expert, in ascending id order, with that
expert's BF16 rows in ascending token order, and returns the expert's
unweighted BF16 output. An expert no token selected is never called.
The component owns dispatch and combination. It checks its inputs and
every provider output (type, dtype, shape, device) with errors naming
the argument; the id range and per-row uniqueness are caller
preconditions and are not scanned on the device.
The arithmetic is one owned op, BeamDispatchCombine, written with
ordinary torch CUDA operations: each weighted route is the FP32 product
of the BF16 expert output and its FP32 weight, rounded once to BF16 and
added to its token's FP32 accumulator in ascending expert order; the
sum is multiplied by the routed scaling factor 2.5 once and rounded to
BF16. The selected expert list and each expert's selection are host
metadata; gathers, products, roundings, sums and the scaling run on the
GPU.
The tests run the public call at seeds {17, 41} x tokens {1, 17, 129}
with controlled affine experts whose outputs depend on the expert id
and the token values, and require bit-exact agreement with a plain-torch
golden of the combination. Focused fixtures cover experts 0 and 255
shared by every token with unequal, unnormalized weights, a signed
cancellation witness on which each rounding step and the single scale
are visible, zero-weight routes, all-zero weights, all-zero expert
outputs and token-row permutation. They also check rejected inputs and
provider outputs, that the inputs are not modified, and that the call
runs on the GPU. The op's test drives its cells bit-exactly against its
own reference.
The tests are collected by the existing l0_b300 pre-merge entry for
unittest/_torch/modeling_v2: 51 tests, about 3 s on one GB300.
To run the tests:
python -m pytest tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
…lingV2
Compose the accepted Beam MoE router, routed SwiGLU expert and
dispatch/combine components into a directly constructed MoE feed-forward
component, and build decoder block 3 (sliding attention, sparse MoE) on
it. Both are components, not routed targets.
- moe_bf16__sm_103__tp1: BeamMoeBf16Sm103Tp1 owns the FP32 router and
256 independently parameterized BF16 experts, checks all 770 tensors
before loading any of them, and calls the accepted dispatch/combine
with each expert as the provider.
- sparse_block_bf16__sm_103__tp1: BeamSparseBlockBf16Sm103Tp1 runs the
input RMSNorm, the sliding-window attention, the four sandwich norms,
the MoE and block 3's depth-scaled residual adds. Its KV cache slot is
the KVCacheManagerV2's local index for logical layer 3
(layer_offsets[3], selected with a layer_mask); a manager that does
not hold layer 3 is rejected before any kernel runs.
- _sliding_attention.py: the attention sublayer, step-argument
projection and call checks shared by the dense and sparse blocks. The
dense block delegates to it and resolves its cache slot the same way.
- Promote beam_neox_rope, beam_output_gate and beam_depth_residual, with
their tests, to the shared catalog, and add block 3's residual cells.
Tests: generated-weight GPU tests for both components cover
construction, loading, input and cache-metadata contracts, the six
fixed cases against plain-torch references, designed routing fixtures,
causal-prefix independence, residual coefficients and fault
sensitivity. The dense block and moved operator tests are updated.
To run the tests:
python -m pytest tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1 \
tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Compose embeddings, three dense and three sparse layers, final normalization, an untied output head and softcap using generated weights. Generalize shared attention for logical layer slots and full-context NoPE query scaling. Validated eager TP1 on GB300 for seeds 17/41 and lengths 1/17/129: 80 public checks and 354 component regressions passed, with independent review and QA acceptance. Scope is six-layer fresh-context correctness; cache reuse, decoding, real weights and full-depth inference remain future work. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Drive the six-layer ModelingV2 component with KVCacheManagerV2 and TorchSampler for eager single-GPU greedy generation and shared-prefix reuse. Validate request history and model-weight identity before cache access, and add generated-data lifecycle and cache-safety tests. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Add bounded streaming checkpoint loading, rank-local ownership, and representative distributed block execution on four GPUs. Include generated-data loader/layout tests and a model-specific RMSNorm operator. This qualifies checkpoint-derived BF16 execution and full parameter residency; full-depth inference and native quantized execution remain outside this component. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Load packed NVFP4 expert weights and scales into rank-local storage and execute representative sparse blocks with TRT-LLM CUTLASS kernels. Preserve BF16 excluded modules and FP32 routers, and add generated-data format, scale, sharding and rejection tests. Qualifies full 52-layer parameter residency and representative eager component execution on four GB300 GPUs. Full-depth inference, generation, task accuracy, CUDA graphs and performance qualification remain future work. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Integrate all 52 layers with the ModelingV2 target, native checkpoint loading, KV cache management, TorchSampler and text chat-completions serving. Support eager TP4 execution and native request batching; add parser/routing and batch metadata coverage, a launch guide and CPU CI registration. Complete the registered 10-shot GSM8K evaluation: 1263/1319 (95.75%), compared with a recorded BF16 mean of 95.53%. Batching numerical qualification remains partial; serial-equivalent trajectories and optimized throughput are not claimed. Independent reviewer and QA accepted the scoped functional baseline. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Merge upstream main while preserving Beam and upstream catalog cases, attention test helpers, and CPU CI registration. Adapt Beam routing to the engine's effective KV-cache argument and configuration keywords, rejecting unsupported quantized KV formats before target construction. Add a Beam deployment guide to Model Recipes with the four-GB300 launcher, chat API example, capacity guidance, measured accuracy and scope limits. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Keep rank-local metadata, streaming and cleanup failures inside the all-rank load agreement. Validate dense-block weights before copying and reject execution until a complete load succeeds. Reject unsupported launcher token budgets and tensor parallelism in the unsharded loader. Make the fallback-selection test independent of the host GPU and skip CUDA metadata setup when no device is available. Add CPU regression coverage for these contracts and Beam config loading. Validation: 31 focused CPU cases passed with native imports isolated; repository pre-commit checks passed. Native GPU execution was not rerun. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
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.
Description
Continuation of #20499 for Beam inference and serving with StairCase/ModelingV2. This draft starts with DFlash2 speculative decoding on the full 52-layer NVFP4 target and will collect the remaining feature tiers listed below as they are completed.
Dependency: #20499 is still open. This branch includes its latest head,
96ce0eca1be9f3134ee3fa3a4e3e9d99a0709ca9, including the CodeRabbit loading, serving-validation, configuration, and test fixes. Until that PR merges, GitHub's main-based diff also includes its baseline changes. The incremental DFlash2 change is commit57fed0768ec3c932174342e7378e73d1c5d24e86(compare against #20499's current head).Included now
--draft-checkpointlaunch support, generated-input regressions, and documented invocation and limits.Qualified configuration: 4×GB300, attention TP4, MoE TP4/EP1, PP1, NVFP4 target experts, BF16 target and drafter KV, eager execution without overlap or chunked prefill. Bounded full-model checks cover batches 1/4, prompts through 4097 tokens, and up to 64 generated tokens. This does not claim speculative GSM8K accuracy or production throughput.
Remaining roadmap — not implemented or qualified by this commit
Test Coverage
Validation limits
QA reused saved full-model GPU evidence as permitted by the task; its own fresh real-checkpoint rerun did not execute before allocation expiry. Independent ordinary/speculative runs can diverge at near ties as GEMM/collective shapes change; exact verification is established against each step's actual target logits. Some serving-versus-direct output differences remain unlocalized. The bounded batch-four timing case queued its fourth long request because of the configured KV capacity. Graphs, overlap, FP8 KV, and disaggregated serving are not yet qualified.
Only implementation, generated-input tests, and documentation are published. No checkpoint weights, private captures, credentials, or workflow artifacts are included.
PR Checklist