Repository navigation
[TRTLLM-17362][feat] Integrate Beam with standard serving and public evaluation - #21473
Mgluhovskoi wants to merge 49 commits into
Conversation
Provide dense and MoE components, cached BF16 small-model execution and the 52-layer NVFP4 TP4 target with BF16 KV. Keep reviewed loading checks and opt-in diagnostics. Include the shared paged-context capability dependency. Exclude the test suite, standard serving registration, speculative decoding, FP8 KV and disaggregation. This is the first boundary in the four-part merge-back sequence. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Add operator, component, layout, cached-generation and diagnostic coverage with shared fixtures. Preserve the runtime-polish simplification. Keep serving and optimization-specific tests for their corresponding merge-back commits; retain the existing single-GPU CI selection. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
…evaluation Register Beam's checkpoint loader, concrete tokenizer and channel parser for LLM() and trtllm-serve. Add baseline serving documentation, parser/loading checks and the public resumable GSM8K evaluator. Keep BF16 KV and eager execution; leave optimized recipes and their scheduled accuracy regression to the next commit. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
|
/bot run |
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
|
/bot run |
|
PR_Github #83521 [ run ] triggered by Bot. Commit: |
|
PR_Github #83515 [ run ] completed with state |
Keep the six TP1 helpers and generation controller out of the runtime wheel. Remove unused TP4 intermediate-capture hooks and observe expert intermediates through test-scoped instrumentation instead. Update scoped imports, documentation and explicit test selectors without broadening CI. Host contracts, wheel contents, static CI selection and pre-commit pass. Native GPU and four-rank qualification remain pending. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
|
/bot run |
Remove the unused runtime weight_state/check_weight_state API and its tensor fingerprint helper. Preserve mutation/replacement checks as private helpers beside the checkpoint-layout test and remove the stale runtime documentation. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
|
/bot run |
|
PR_Github #83527 [ run ] triggered by Bot. Commit: |
|
PR_Github #83521 [ run ] completed with state |
|
PR_Github #83532 [ run ] triggered by Bot. Commit: |
|
PR_Github #83527 [ run ] completed with state |
Use descriptive checkout-only component names, including six_layer_model for prefill and decoding. Move shared BF16 TP4 layout/loading/component files into family _tp4. Name the operator beam_dispatch_combine and its checked wrapper dispatch_and_combine. Update imports and test selectors; preserve canonical model and routed-target identifiers. Only files introduced by the Beam PRs are renamed. Wheel contents, static CI selection, host contracts and pre-commit pass; native CI remains pending. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
|
/bot run |
|
PR_Github #83535 [ run ] triggered by Bot. Commit: |
|
PR_Github #83532 [ run ] completed with state |
|
I re-checked the Blockers from my review of 2026-11-04 at cd11961. None of the 14 open Blockers has a reply or a code change since 2ab3329. At cd11961, three of them still fail pre-merge CI:
Resolved or accepted:
Still open, so my request for changes stands:
|
|
PR_Github #83535 [ run ] completed with state
|
| flashinfer_rmsnorm, | ||
| ) | ||
| from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( | ||
| _attention, |
There was a problem hiding this comment.
Could you update the tests to import _attention and dispatch_combine from models.beam after moving those modules? This import still targets full_model_nvfp4__sm_103__tp4, which no longer contains _attention, so the dense-block suite selected by l0_b300.yml fails during collection. The MoE suite has the same stale dispatch_combine import.
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
|
/bot run |
|
PR_Github #83584 [ run ] triggered by Bot. Commit: |
|
PR_Github #83584 [ run ] completed with state
|
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Keep stock checkpoint registration independent of ModelingV2 targets. Preserve the target loader import as a compatibility alias, reset preload status on every attempt, and retain configured HF loaders and mappers across worker serialization. Add CPU lifecycle and serialization coverage to L0. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Warn when auto routing has no built-in fallback, aggregate eager-engine and DFlash refusal reasons, and read the token budget without importing target model code. Preserve accepted BF16/FP8-cache and DFlash routes. Add CPU admission coverage to L0. Completes propagation of the review fix from NVIDIA#21473 to NVIDIA#21474 for NVIDIA#21471 discussion r4470150943. Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
|
I re-checked the Blockers from my review of 2026-11-04 at 37f775f. At 37f775f the three Blockers that failed pre-merge CI (docs build, stock-import test, api_stability reference) are fixed and the guide config gets past KV-cache estimation, but the other eleven Blockers have no reply and their code is unchanged since 2ab3329. Resolved or accepted:
Still open, so my request for changes stands:
|
Description
Integrate the Beam 52-layer NVFP4 baseline with standard
LLM()andtrtllm-servepaths: ModelingV2 routing and engine adapter, registeredbeam_nvfp4checkpoint loader, tokenizer, Harmony parser, serving documentation and the public GSM8K evaluator.Part 3/4 of TRTLLM-17024, implementing TRTLLM-17362. Depends on #21471 and #21472.
The runnable target is introduced here alongside its engine-admission tests and named
TestModelingV2BeamFullModelNvfp4Sm103Tp4accuracy gate. Earlier parts contain directly constructed components only. Component corrections are inherited, and the loader guard accepts the documentedcheckpoint_format="beam_nvfp4"path while rejecting incompatible explicit loaders.The baseline is TP4 on four GB300s, NVFP4 experts, BF16 KV, eager execution and overlap disabled. Speculative decoding, FP8 KV, disaggregated serving and other optimizations are outside this baseline.
Review order
#21471 → #21472 → #21473 → #21474. All target main; later PRs are cumulative.
Incremental part-3 diff
Validation
At
62a184ba5383e9449f6ee906b13517bf045b6215:The named gate starts
trtllm-servewithmodeling_v2: requireand runs the pinned public first-100 GSM8K protocol: at least 95 correct, with execution/parsing/truncation failures rejected. It is a manual four-GB300 qualification with a 45-minute cap, not an L0 or full-dataset CI run. Authorized checkpoint and dataset inputs must be staged externally; missing inputs fail.The gate has not been executed on a matching GPU build for this corrected extraction. Native accuracy/timing qualification is required before scheduling it in LR/nightly or claiming this baseline qualified. Remaining human review, matching-build GPU validation and required CI/code-owner approval remain outstanding. No checkpoint weights or private evidence are included.
Dev Engineer Review
Adds Beam support across the PyTorch backend and the ModelingV2 serving path. The changes include Beam model components and operators, TP4 NVFP4 checkpoint loading and routing, tokenizer and Harmony parser support, serving documentation, and a pinned GSM8K evaluator.
The engine route accepts a narrow four-rank SM 10.3 NVFP4 configuration and refuses unsupported serving options. The documented baseline uses BF16 KV cache and disables overlap scheduling. Speculative decoding, FP8 KV cache, and disaggregated serving are outside its scope.
The evaluator supports managed-server and existing-endpoint modes, resumable records, model identity checks, and explicit failure verdicts. The manual gate requires at least 95 correct answers on the pinned first 100 GSM8K questions. The supplied context says matching-build GPU qualification remains outstanding. It also reports one successful PR job and one failed merge-request pipeline; investigation and a rerun remain outstanding.
QA Engineer Review
The change adds tests for Beam operators and components, checkpoint layouts and loading, routing and engine contracts, diagnostic behavior, Harmony parsing, evaluator protocol and records, and the accuracy gate. The
l0_cpu.ymllist includes component-contract, engine-contract, baseline-gate, diagnostic-mode, loading-contract, and serving-parser tests. It does not includetests/unittest/tools/test_nvfp4_serving_gsm8k_first100.py. The manual GSM8K gate is listed inmanual_gb300_beam_baseline_gsm8k_first100.ymlfor the four-GB300 profile, with a 45-second test-list timeout.Coverage verdict: needs follow-up. CPU contract tests cover key evaluator and serving paths, but they do not establish GPU compatibility. Run the manual accuracy gate on a matching GPU build and investigate the reported failed merge-request pipeline.
Per-File QA Perspective
Documentation and evaluator
docs/source/deployment-guide/deployment-guide-for-beam-on-trtllm.md: Verify the documented serving configuration and manual accuracy criteria against runtime behavior.docs/source/deployment-guide/index.rst: Adds the Beam guide to the deployment documentation tree.examples/llm-api/eval_gsm8k.py: Verify CLI validation, endpoint and managed-server modes, runner arguments, and exit codes.examples/llm-api/gsm8k_eval/README.md: Check that usage and failure documentation matches evaluator behavior.examples/llm-api/gsm8k_eval/__init__.py: Adds package documentation; no runtime behavior change is described.examples/llm-api/gsm8k_eval/manifest.json: Verify the pinned dataset, question and exemplar selection, and recorded hashes.examples/llm-api/gsm8k_eval/protocol.py: Verify manifest validation, prompt construction, selection, and scoring.examples/llm-api/gsm8k_eval/records.py: Verify reply classification, ledger validation and resume behavior, and verdict rules.examples/llm-api/gsm8k_eval/runner.py: Verify endpoint and managed-server execution, model consistency, resume identity, and interruption handling.examples/llm-api/gsm8k_eval/server.py: Verify server startup, address reporting, process cleanup, and optional GPU memory polling.ModelingV2 operators and Beam model
tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py: Verify the resolver documentation reflects its import behavior.tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/beam_logit_softcap.py,flashinfer_silu_and_mul.py, andswiglu_fp32.py: Verify operator formulas, dtype and shape constraints, and new SM 103 cases.tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_neox_rope.py,beam_nope_query_scale.py,beam_output_gate.py, andthop_attention.py: Verify RoPE and query scaling, output gating, attention windows, and FMHA case configurations.tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py,catalog/moe/beam_dispatch_combine.py,beam_grouped_topk_router.py,catalog/norm/beam_rmsnorm_fp32.py,flashinfer_rmsnorm.py, andcatalog/residual/beam_depth_residual.py: Verify projection coverage, routing and accumulation order, normalization, and residual scaling.tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/__init__.py: Adds package metadata; no runtime behavior change is described.tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md: Check documented target constraints and unsupported serving features against implementation.tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/__init__.pyand the Beam component package__init__.pyfiles: These add package metadata; no runtime behavior change is described.tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/block.py,moe_bf16__sm_103__tp1/moe.py,moe_expert_bf16__sm_103__tp1/expert.py,moe_router_bf16__sm_103__tp1/router.py, andsparse_block_bf16__sm_103__tp1/block.py: Verify supported configurations, weight-loading guards, attention and routing behavior, and CUDA device checks.tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/model.pyandgeneration.py: Verify cache lifecycle, prefix reuse, weight-change detection, and greedy generation.tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/_attention.py,component.py,component_bf16.py,dispatch_combine.py,expert.py,layout.py,layout_bf16.py,loader.py, andloader_bf16.py: Verify attention metadata and cache checks, rank-local checkpoint reads, quantization layout, collective behavior, and load-failure handling.tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/model.pyandmodeling.py: Verify full-model execution, prefill/decode dispatch, request restrictions, and logit selection.tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/weights.pyandmodels/beam/routing.py: Verify checkpoint-loader integration and engine admission/refusal behavior.tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py: Verify importing the package loads Beam routing without unintended registration or import effects.tensorrt_llm/_torch/models/checkpoints/__init__.py,tensorrt_llm/_torch/models/checkpoints/beam/checkpoint_loader.py, andtokenizer.py: Verify public exports,beam_nvfp4registration, and tokenizer loading.tensorrt_llm/_torch/models/checkpoints/beam/__init__.py: Adds package metadata; no runtime behavior change is described.tensorrt_llm/_torch/models/modeling_beam.py: Verify model computation, configuration validation, and strict checkpoint mapping and loading.tensorrt_llm/_torch/pyexecutor/config_utils.py: Verify Beam config construction preserves the checkpoint path and supplied settings.tensorrt_llm/llmapi/reasoning_parser.py: Verify complete and streamed Harmony channel parsing, including termination markers.Tests and test lists
tests/integration/defs/accuracy/modeling_v2/test_beam.py: Adds the manual first-100 GSM8K gate. It is listed in the manual GB300 test list.tests/integration/test_lists/test-db/manual_gb300_beam_baseline_gsm8k_first100.yml: Adds the manual gate selector for the four-GB300 profile. Verify its selectors and timeout against the test.tests/integration/test_lists/test-db/l0_cpu.yml: Adds CPU component-contract, engine-contract, baseline-gate, diagnostic, loading, and serving tests. The evaluator protocol/records test is not listed.tests/unittest/tools/test_nvfp4_serving_gsm8k_first100.py: Covers protocol, scoring, dataset checks, reply handling, and ledger behavior with synthetic data. It is not listed inl0_cpu.yml.tests/unittest/tools/test_beam_baseline_accuracy_gate.py: Covers gate arguments, failure propagation, and checkpoint-path validation. It is listed inl0_cpu.yml.tests/unittest/_torch/modeling_v2/models/beam/test_component_contracts.pyandtest_engine_contracts.py: Cover cache and checkpoint contracts, request validation, engine-option refusal, and checkpoint-loader routing. Both are listed inl0_cpu.yml.tests/unittest/_torch/modeling_v2/models/beam/test_loading_contracts.pyandtest_diagnostic_mode.py: Cover loading failures, weight reload behavior, and rank agreement. Both are listed inl0_cpu.yml.tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_serving.py: Covers parser registration and pretrained config loading. It is listed inl0_cpu.yml.tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_batching.py: Covers batched attention metadata and logit-row selection. No test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/test_checkpoint_layout.pyandcheckpoint_layout_nvfp4__sm_103__tp4/test_checkpoint_layout_nvfp4.py: Cover layout partitioning, checkpoint metadata, rank slices, and CUDA loading paths. No test-list entries are shown.tests/unittest/_torch/modeling_v2/models/beam/_beam_component_support.py: Adds shared test fixtures and evidence helpers; it is not a test module listed in a test list.tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/_dense_block_reference.py: Adds a Torch reference for dense-block calculations; no test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/test_dense_block.py: Covers dense-block reference comparisons, restrictions, execution paths, and KV-cache writes. No test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1/_moe_reference.py: Adds a Torch reference for MoE routing and expert computations; no test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1/test_moe.py: Covers MoE routing, numerical behavior, configuration, weights, and execution paths. No test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/_dispatch_combine_cases.py: Adds generated fixtures and a golden dispatch/combine implementation; no test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/moe/test_modeling_v2_beam_dispatch_combine.pyandtest_dispatch_combine.py: Cover operator certification and component routing, output ordering, and input immutability. No test-list entries are shown.tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/_expert_cases.py: Adds numerical controls, fixtures, and a golden expert implementation; no test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/activation/test_modeling_v2_swiglu_fp32.py: Covers SwiGLU reference and precision behavior, with validation guards. No test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/test_expert.py: Covers expert loading, validation, numerical stages, and execution behavior. No test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/_router_cases.py: Adds router fixtures and comparison helpers; no test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/moe/test_modeling_v2_beam_grouped_topk_router.pyandtest_router.py: Cover router reference comparisons, routing fixtures, validation, and numerical controls. No test-list entries are shown.tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_beam_logit_softcap.py: Covers softcap operator behavior and validation. No test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_cached_generation.py: Covers generation lifecycle, cache writes, prefix reuse, and weight-change detection. No test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_small_model_prefill.py: Covers prefill validation, execution order, attention behavior, and request isolation. No test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/_sparse_block_reference.py: Adds a Torch reference for sparse-block calculations; no test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/test_sparse_block.py: Covers sparse-block output, routing, configuration, and cache behavior. No test-list entry is shown.tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_neox_rope.py,test_modeling_v2_beam_nope_query_scale.py, andtest_modeling_v2_beam_output_gate.py: Cover operator references, shape and dtype behavior, precision, and validation guards. No test-list entries are shown.tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py: Extends attention reference and paged-context coverage for Beam sliding and full attention. No test-list entry is shown.tests/unittest/_torch/modeling_v2/residual/test_modeling_v2_beam_depth_residual.py: Covers residual coefficients, reference output, and validation guards. No test-list entry is shown.tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/norm/test_modeling_v2_beam_rmsnorm_fp32.py: Covers RMSNorm reference precision, output properties, and validation. No test-list entry is shown.