[None][feat] Extend Beam ModelingV2 inference and serving

docs/source/deployment-guide/deployment-guide-for-beam-on-trtllm.md

@@ -0,0 +1,131 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+-->
+
+# Deployment Guide for Beam on TensorRT LLM
+
+This guide serves the complete 52-layer Beam model with NVFP4 weights on **four
+GB300 GPUs on one node**, using TensorRT LLM's experimental StairCase/ModelingV2
+implementation. It provides streaming and non-streaming text chat completions
+through the native OpenAI-compatible server.
+
+This is a functional deployment baseline. The expert GEMMs use native NVFP4
+kernels, but expert dispatch currently uses a host-driven loop; this recipe is
+not a throughput-optimized configuration.
+
+## Prerequisites
+
+- Four NVIDIA GB300 GPUs (SM103), visible to the same Linux environment, with
+  working GPU peer communication. The full-model target requires TP4; single-GPU
+  execution and other parallel layouts are not supported by this recipe.
+- A TensorRT LLM installation containing the Beam ModelingV2 target and matching
+  native libraries. Follow the [source build instructions](../installation/build-from-source.md)
+  for the checkout containing this guide. An older release image may not contain
+  this experimental target.
+- An authorized, compatible Beam NVFP4 checkpoint directory, accessible to every
+  rank, containing `config.json`, the tokenizer files, `chat_template.jinja`,
+  `model.safetensors.index.json`, and all indexed weight shards. Preserve the
+  checkpoint's quantization metadata and scales. The launcher does not download
+  or convert weights.
+
+## Deployment configuration
+
+| Setting | Configuration |
+|---|---|
+| Model | All 52 layers, with checkpoint-defined sliding-window and full attention |
+| Parallelism | Attention TP4, MoE TP4/EP1, PP1; attention data parallelism disabled |
+| Precision | Packed NVFP4 experts with native W4A4 GEMMs; BF16 excluded modules and activations; FP32 router |
+| Attention backend | TRTLLM |
+| KV cache | BF16, `KVCacheManagerV2`, full-block prefix reuse enabled, partial-block reuse disabled |
+| Sampling | Native `TorchSampler`; greedy decoding is the evaluated policy |
+| Execution | Eager; CUDA graphs, overlap scheduling, and chunked prefill disabled |
+| Server | Native OpenAI-compatible server, one frontend, text-only |
+
+The Beam launcher configures the checkpoint loader, tokenizer, reasoning parser,
+ModelingV2 routing, and native server together. Use this launcher for the recipe;
+plain `trtllm-serve` with default arguments does not establish the same setup.
+There is no separate serving YAML required for this launch.
+
+## Launch the server
+
+Run one launcher process with access to all four GPUs; it starts the TP4 workers.
+Replace `/path/to/checkpoint` with your checkpoint directory.
+
+```bash
+python -m tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.serving \
+  --checkpoint /path/to/checkpoint \
+  --served-model-name beam-nvfp4 \
+  --host 127.0.0.1 --port 8000 \
+  --max-seq-len 69632 \
+  --max-num-tokens 8192 \
+  --max-batch-size 1
+```
+
+The server uses the checkpoint's `chat_template.jinja` by default. Use
+`--chat-template /path/to/chat_template.jinja` only when supplying a compatible
+template explicitly. The example binds to loopback; choose a reachable bind
+address when deploying behind your service's access controls.
+
+### Capacity and batching
+
+- `--max-seq-len` bounds prompt plus generated tokens. The configured capacity
+  is not a claim of evaluated accuracy at that context length.
+- `--max-num-tokens` bounds tokens processed in an engine step. With chunked
+  prefill disabled, the prompt must fit this budget. This target accepts at most
+  8,192 new token rows per model call.
+- `--max-batch-size 1` is the conservative starting configuration. Increasing it
+  enables native scheduler batching and requires a corresponding KV memory budget.
+- `--kv-cache-max-tokens` sets the shared cache capacity; by default it is twice
+  `--max-seq-len`. Capacity must account for all concurrently active requests.
+  The GSM8K evaluation used a batch limit of 128 and a cache token budget of
+  1,048,576, which are workload-specific settings, not general capacity guarantees.
+- `--iter-stats` enables native iteration statistics exposed through `/metrics`.
+
+## Send a chat request
+
+After the server is ready:
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "beam-nvfp4",
+    "messages": [{"role": "user", "content": "What is 12 plus 7?"}],
+    "temperature": 0,
+    "max_tokens": 4096,
+    "stream": false,
+    "chat_template_kwargs": {"reasoning_effort": 128}
+  }'
+```
+
+The response's `choices[0].message.content` contains the final answer;
+`choices[0].message.reasoning_content` contains analysis-channel text. The parser
+removes the channel delimiters. Set `"stream": true` and add curl's `-N` option
+for streaming; read `content` and `reasoning_content` from each choice's `delta`.
+
+`max_tokens` includes reasoning tokens. If the response reaches its token limit
+while reasoning, it may have no final answer; increase the completion allowance
+within the configured sequence capacity. The example illustrates the API and
+does not guarantee completion within 4,096 tokens.
+
+## Accuracy and limitations
+
+The initial full-model evaluation scored **1,263/1,319 (95.75%) on GSM8K**, using
+10-shot prompting, greedy decoding, reasoning effort 128, and a 65,536-token
+completion allowance. All questions completed without execution errors or
+truncated completions. The recorded comparison is a **95.53% BF16 mean** over
+three runs, not a same-precision NVFP4 reference. Differences in precision,
+runtime, and batching prevent interpreting the small score increase as a benefit
+of quantization.
+
+Greedy outputs can differ across batch sizes and between fresh and reused-cache
+execution. The evaluation does not establish exact token parity for these modes.
+Long-context accuracy, performance targets, and a checkpoint-backed CI regression
+remain separate qualification work.
+
+FP8 KV cache, speculative decoding, CUDA graphs, overlap scheduling, chunked
+prefill, tool calling, other GPU architectures/layouts, and disaggregated serving
+are not qualified by this recipe. The model's implementation and generated-input
+regression coverage are described in the
+[Beam source README](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md).

docs/source/deployment-guide/index.rst

@@ -30,6 +30,7 @@ The deployment guides below provide more detailed instructions for serving speci
    :maxdepth: 1
    :name: Deployment Guides
 
+   deployment-guide-for-beam-on-trtllm.md
    deployment-guide-for-nemotron-3-on-trtllm.md
    deployment-guide-for-gpt-oss-on-trtllm.md
    deployment-guide-for-qwen3-on-trtllm.md

tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py

@@ -51,6 +51,7 @@
 
 # architectures[0] -> routing module, relative to this package.
 MODELING_V2_ROUTERS = {
+    "BeamForCausalLM": "models.beam.routing",
     "GptOssForCausalLM": "models.gpt_oss.routing",
     "DeepseekV3ForCausalLM": "models.deepseek_v3.routing",
     # Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the

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

@@ -23,7 +23,8 @@ class FlashinferSiluAndMul(OpWrapper):
 
     # 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.
+    # -- at decode and prefill token counts. The Beam dense block component
+    # calls it at its 20480 intermediate and three context lengths.
     CELLS: tuple[Cell, ...] = (
         Cell(
             why="decode-shaped, 2 x moe_intermediate (2048)",
@@ -37,6 +38,13 @@ class FlashinferSiluAndMul(OpWrapper):
             why="the dense layers' wider intermediate, 2 x 18432",
             spec=dict(shape=(16, 36864), dtype=torch.bfloat16),
         ),
+        *(
+            Cell(
+                why=f"beam dense block, {tokens} tokens: 2 x 20480",
+                spec=dict(shape=(tokens, 40960), dtype=torch.bfloat16),
+            )
+            for tokens in (1, 17, 129)
+        ),
         Cell(
             why="deepseek default at attention TP8: the shared expert's gate+up, 2 x 2048/8",
             spec=dict(shape=(16, 512), dtype=torch.bfloat16),

tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_neox_rope.py

@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""NeoX half-rotation RoPE on per-head q and k, in fp32 with one rounding per tensor."""
+
+from __future__ import annotations
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: Block 0's sliding-attention layer: 48 q heads and 8 kv heads of 128, theta 10000, factor 1.
+#: Block 3's sliding-attention layer has the same geometry and RoPE parameters.
+_BLOCK0 = dict(
+    num_heads=48, num_kv_heads=8, head_dim=128, theta=10000.0, factor=1.0, dtype=torch.bfloat16
+)
+
+
+def _cos_sin(
+    positions: torch.Tensor, head_dim: int, theta: float, factor: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """fp32 `[T, 1, head_dim]` cos and sin, both halves carrying the same angles."""
+    exponent = torch.arange(0, head_dim, 2, dtype=torch.float32, device=positions.device)
+    inv_freq = 1.0 / (theta ** (exponent / head_dim)) / factor
+    freqs = positions.float().unsqueeze(-1) * inv_freq
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos().unsqueeze(-2), emb.sin().unsqueeze(-2)
+
+
+def _rotate(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    x32 = x.float()
+    first, second = x32.chunk(2, dim=-1)
+    rotated = torch.cat((-second, first), dim=-1)
+    return (x32 * cos + rotated * sin).to(x.dtype)
+
+
+class BeamNeoxRope(OpWrapper):
+    """Rotate `q` `[T, Hq, D]` and `k` `[T, Hkv, D]` by their token's position.
+
+    Frequencies are `1 / theta ** (2i / D) / factor` for `i < D / 2`; angles are
+    `position * frequency` in fp32, and the two halves of the head dimension rotate
+    together -- `x * cos + rotate_half(x) * sin`, NeoX order, not interleaved pairs.
+    Each output is rounded once to its input's dtype. `q` and `k` arrive already
+    normalized and rounded by their per-head RMSNorm; this op does not normalize, and
+    `v` never passes through it. Returns two new tensors.
+    """
+
+    ARCHS = frozenset({Arch.SM_103})
+
+    # The three fresh-context lengths the dense block runs: a single token (position 0,
+    # the identity rotation), a short and a long context.
+    CELLS: tuple[Cell, ...] = (
+        Cell(
+            why="block 0, 1 token: position 0 is the identity rotation",
+            spec=dict(_BLOCK0, tokens=1),
+        ),
+        Cell(why="block 0, 17 tokens: positions 0..16", spec=dict(_BLOCK0, tokens=17)),
+        Cell(
+            why="block 0, 129 tokens: the longest context the block accepts",
+            spec=dict(_BLOCK0, tokens=129),
+        ),
+    )
+
+    note = """
+    Native torch, not a fused kernel: the source model rounds the per-head norm
+    output to the activation dtype before promoting it again for the rotation, and
+    no fused norm-plus-RoPE op in this tree has that intermediate rounding (the
+    catalog's fused_qk_norm_rope keeps the normalized values in fp32 through the
+    rotation). `q` and `k` share one cos/sin evaluation.
+
+    `positions` is per token, not per sequence: a fresh context passes
+    `arange(T)`. No scaling other than the linear `factor` is modelled.
+    """
+
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        positions: torch.Tensor,
+        theta: float,
+        factor: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        cos, sin = _cos_sin(positions, q.shape[-1], theta, factor)
+        return _rotate(q, cos, sin), _rotate(k, cos, sin)
+
+    def reference(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        positions: torch.Tensor,
+        theta: float,
+        factor: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """The rotation written per half: `(x1 c - x2 s, x2 c + x1 s)`, fp32, one cast."""
+        half = q.shape[-1] // 2
+        index = torch.arange(half, dtype=torch.float32, device=positions.device)
+        inv_freq = 1.0 / (theta ** (2.0 * index / q.shape[-1])) / factor
+        angle = positions.float()[:, None] * inv_freq[None, :]
+        c, s = angle.cos()[:, None, :], angle.sin()[:, None, :]
+
+        def rotate(x: torch.Tensor) -> torch.Tensor:
+            x1, x2 = x.float()[..., :half], x.float()[..., half:]
+            return torch.cat((x1 * c - x2 * s, x2 * c + x1 * s), dim=-1).to(x.dtype)
+
+        return rotate(q), rotate(k)
+
+    def is_valid(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        positions: torch.Tensor,
+        theta: float,
+        factor: float,
+    ) -> None:
+        # A [1] or [B, 1] positions tensor broadcasts across every token and rotates
+        # them all by one angle; a float tensor is accepted and rotates by fractional
+        # positions. Neither raises.
+        assert q.dim() == 3 and k.dim() == 3, "q and k must be [T, heads, head_dim]"
+        assert positions.shape == (q.shape[0],) and k.shape[0] == q.shape[0], (
+            f"positions must be [T] with T = {q.shape[0]} tokens; got {tuple(positions.shape)} "
+            f"for k of {k.shape[0]} tokens"
+        )
+        assert not positions.is_floating_point(), "positions must be integer token positions"
+        assert q.shape[-1] == k.shape[-1] and q.shape[-1] % 2 == 0, (
+            "q and k must share an even head_dim"
+        )
+
+
+beam_neox_rope = BeamNeoxRope()

tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_nope_query_scale.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Position-dependent query scale of a full-attention layer without position embedding."""
+
+from __future__ import annotations
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: Block 5's full-attention layer: 48 q heads of 128, bf16.
+_BLOCK5 = dict(num_heads=48, head_dim=128, dtype=torch.bfloat16)
+
+
+class BeamNopeQueryScale(OpWrapper):
+    """Scale `q` `[T, H, D]` per token: `q * exp(0.1 * log1p((position + 1) / 2048))`.
+
+    The scale is computed in fp32 from the token's position, the product is the
+    activation-dtype `q` promoted against that fp32 scale, and the result is rounded
+    once to `q`'s dtype. It replaces the rotation of a full-attention layer that has no
+    position embedding: it applies to the normalized `q` only, never to `k` or `v`,
+    and it is not the attention's `1 / sqrt(D)` softmax scale. Returns a new tensor.
+    """
+
+    ARCHS = frozenset({Arch.SM_103})
+
+    CELLS: tuple[Cell, ...] = (
+        Cell(why="block 5, 1 token: position 0 alone", spec=dict(_BLOCK5, tokens=1)),
+        Cell(why="block 5, 17 tokens: positions 0..16", spec=dict(_BLOCK5, tokens=17)),
+        Cell(
+            why="block 5, 129 tokens: the longest context the block accepts",
+            spec=dict(_BLOCK5, tokens=129),
+        ),
+    )
+
+    note = """
+    The scale grows from 1.00005 at position 0 to 1.0063 at position 128, so most
+    bf16 query values are unchanged by it at these lengths: only values close to a
+    rounding boundary move, and position 0 moves almost none. A test that has to see
+    the scale therefore compares the rounded result exactly at nonzero positions.
+    """
+
+    def __call__(self, q: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        scale = torch.exp(0.1 * torch.log1p((positions.float() + 1) / 2048))
+        return (q * scale[:, None, None]).to(q.dtype)
+
+    def reference(self, q: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        """Each step in fp32, the product of the widened `q` and the scale, one cast."""
+        position = positions.to(torch.float32)
+        ratio = torch.div(position + 1.0, 2048.0)
+        scale = torch.exp(torch.mul(torch.log1p(ratio), 0.1))
+        return (q.to(torch.float32) * scale.view(-1, 1, 1)).to(q.dtype)
+
+    def is_valid(self, q: torch.Tensor, positions: torch.Tensor) -> None:
+        # A [1] positions tensor broadcasts one scale over every token, and a float
+        # tensor scales by fractional positions. Neither raises.
+        assert q.dim() == 3, "q must be [T, heads, head_dim]"
+        assert positions.shape == (q.shape[0],), (
+            f"positions must be [T] with T = {q.shape[0]} tokens, got {tuple(positions.shape)}"
+        )
+        assert not positions.is_floating_point(), "positions must be integer token positions"
+
+
+beam_nope_query_scale = BeamNopeQueryScale()