[TRTLLM-17360][feat] Add Beam StairCase operators and model components

tensorrt_llm/_torch/_experimental/modeling_v2/README.md

@@ -379,3 +379,11 @@ That boundary may move as models need more customization. Model-specific
 runtime policy can be split from shared services, but a change to common
 runtime behavior still has to work for every other model. How to specialize
 the runtime further is still being designed.
+
+### Beam diagnostic validation
+
+For Beam, `TRTLLM_MODELING_V2_VALIDATE=1` additionally enables repeated per-layer
+metadata checks and per-step, all-rank routing agreement using device-to-host copies
+and SHA-256 hashes. The serving adapter also enables checkpoint digests in this mode.
+This is expensive diagnostic behavior, not only a first-forward contract check.
+Use the same setting on every rank and leave it disabled for performance measurements.

tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py

@@ -80,15 +80,19 @@ def registered_routes() -> Dict[str, Callable[..., type]]:
 #: construction, and a check that can only ever hold is a check whose failure
 #: mode is its own bug taking down a working deployment.
 #:
-#: An environment variable rather than an LLM arg while modeling_v2 is
-#: experimental; it has to reach the ranks, so a test passes it through
+#: Beam additionally enables checkpoint digests, per-layer metadata rechecks and per-step
+#: routing agreement (device-to-host copies, hashing and collectives). Keep it off for
+#: performance measurements; all Beam ranks must use the same setting.
+#:
+#: An environment variable while modeling_v2 is experimental. It must reach all ranks,
+#: so a test passes it through
 #: ``env_overrides``. When this stops being experimental it belongs in the
 #: LLM args, where a caller can set it per engine rather than per process.
 MODELING_V2_VALIDATE_ENV = "TRTLLM_MODELING_V2_VALIDATE"
 
 
 def step_contract_enabled() -> bool:
-    """Whether targets should check their step contract on the first forward.
+    """Whether contract diagnostics are enabled (Beam also checks every step).
 
     Read once per model instance at construction rather than per forward: the
     answer cannot change within a process, and the forward is the hot path.

tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py

@@ -21,9 +21,8 @@ class FlashinferSiluAndMul(OpWrapper):
 
     ARCHS = frozenset({Arch.SM_100, Arch.SM_103})
 
-    # deepseek-r1 is the only shipped caller. It reaches this op from two
-    # widths -- the routed/shared expert intermediate and the dense layers'
-    # -- at decode and prefill token counts.
+    # Exercise decode and prefill across narrow and wide intermediates.
+    # Model-specific shapes live in the test driver.
     CELLS: tuple[Cell, ...] = (
         Cell(
             why="decode-shaped, 2 x moe_intermediate (2048)",

tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py

@@ -20,8 +20,8 @@ class CublasMm(OpWrapper):
 
     ARCHS = frozenset({Arch.SM_100, Arch.SM_103})
 
-    # Both targets use this for every dense projection. This entry's test
-    # carries their three argument shapes that differ in kind rather than in size: fused
+    # Targets use this for dense projections. This entry's test
+    # carries argument shapes that differ in kind rather than in size: fused
     # bias (gpt-oss carries one on qkv, o and the router), no bias (deepseek's
     # projections), and an out_dtype that is not the input's (deepseek's router
     # asks for fp32 logits).
@@ -34,7 +34,7 @@ class CublasMm(OpWrapper):
 
     note = """
     `output_buffer_kind` and `group` are the op's multi-stream and collective
-    arguments. Both shipped targets leave them at their defaults, so neither is
+    arguments. The covered targets leave them at their defaults, so neither is
     a cell axis and neither is certified here.
     """
 

tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Residual entries."""

tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md

@@ -0,0 +1,36 @@
+<!-- SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!-- SPDX-License-Identifier: Apache-2.0 -->
+
+# Beam StairCase components
+
+This directory implements Beam components using the experimental ModelingV2/StairCase
+architecture: family-scoped operators and a 52-layer NVFP4-expert component for attention TP4 and MoE TP4/EP1 on four GB300 GPUs.
+
+The full component implements sliding-window and full attention with a BF16 KV cache.
+Its loader streams caller-provided weights into rank-local storage; no checkpoint is
+included. Router projection and selection use FP32, and packed NVFP4 experts use the
+native GEMM path.
+
+These are directly constructed components, not a registered `LLM()` or `trtllm-serve`
+target. Engine routing, the engine adapter, checkpoint-loader registration and a named
+whole-model accuracy gate belong to the serving integration. Operator certification and
+CPU component-contract tests accompany these components; broader component tests are
+reviewed separately. Speculative decoding, FP8 KV cache, disaggregated serving, CUDA
+graphs and overlap scheduling are outside this component baseline.
+
+Beam-only operators live under `ops/`, with tests in the matching Beam test tree.
+Component-specific cases for shared catalog operators reuse the catalog test drivers,
+including TP1 shapes and the per-rank TP4 projection, attention and expert shapes.
+
+## Checkout-only bring-up helpers
+
+The six TP1 bring-up components and their generation controller live with their tests in
+`tests/unittest/_torch/modeling_v2/models/beam/_components/`. They are not installed in
+the TensorRT-LLM wheel. The runtime does not expose intermediate-tensor capture hooks.
+Normal runtime validation and rank-consistency diagnostics remain available.
+
+The TP4 checkpoint layout, expert provider and full-model execution remain in this
+library. Public serving integration uses `LLM()` and `trtllm-serve`.
+
+Shared BF16 TP4 layout, loading and component base modules live in `_tp4/`; the
+`full_model_nvfp4__sm_103__tp4/` package adds NVFP4 expert storage and execution.

tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""BeamForCausalLM components."""