diff --git a/docs/source/deployment-guide/deployment-guide-for-beam-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-beam-on-trtllm.md new file mode 100644 index 000000000000..9521cf29ddcd --- /dev/null +++ b/docs/source/deployment-guide/deployment-guide-for-beam-on-trtllm.md @@ -0,0 +1,131 @@ + + +# 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). diff --git a/docs/source/deployment-guide/index.rst b/docs/source/deployment-guide/index.rst index b79a7e5d1c0f..29139ec8fce1 100644 --- a/docs/source/deployment-guide/index.rst +++ b/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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index cd2a1169d7bd..4776129b796b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py index 3ffd35cc5311..00269ea10ee2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py +++ b/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), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_neox_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_neox_rope.py new file mode 100644 index 000000000000..4b5b952f3f71 --- /dev/null +++ b/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() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_nope_query_scale.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_nope_query_scale.py new file mode 100644 index 000000000000..033e47f15aa8 --- /dev/null +++ b/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() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_output_gate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_output_gate.py new file mode 100644 index 000000000000..5b87f8713ee1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_output_gate.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-head sigmoid output gate with a per-head bias, in the activation dtype.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Block 0's gate: 48 heads of 128, bf16. Block 3's gate has the same geometry. +_BLOCK0 = dict(num_heads=48, head_dim=128, dtype=torch.bfloat16) + + +class BeamOutputGate(OpWrapper): + """`attention_output * sigmoid(gate + gate_bias)`, all `[T, H, D]` but `gate_bias` `[H]`. + + `gate_bias` is broadcast over tokens and the head dimension. The add, the sigmoid + and the product are separate activation-dtype operations, each rounded to that + dtype, which is what the source model computes; an fp32 chain rounded once is a + different result. Returns a new tensor. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block 0, 1 token", spec=dict(_BLOCK0, tokens=1)), + Cell(why="block 0, 17 tokens", spec=dict(_BLOCK0, tokens=17)), + Cell( + why="block 0, 129 tokens: the longest context the block accepts", + spec=dict(_BLOCK0, tokens=129), + ), + ) + + note = """ + The gate is the second q-width projection of the attention input, not a slice of + q, and it is applied to the attention core's output before the output projection. + Omitting the bias, or using one value for every head, still returns a plausible + gated output; only a head-varying bias makes either mistake visible, which is why + the block's tests use one. + """ + + def __call__( + self, + attention_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor, + ) -> torch.Tensor: + return attention_output * torch.sigmoid(gate + gate_bias.unsqueeze(-1)) + + def reference( + self, + attention_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor, + ) -> torch.Tensor: + """Each step in fp32, rounded to the activation dtype before the next.""" + dtype = attention_output.dtype + biased = (gate.float() + gate_bias.float().unsqueeze(-1)).to(dtype) + weight = torch.sigmoid(biased.float()).to(dtype) + return (attention_output.float() * weight.float()).to(dtype) + + def is_valid( + self, + attention_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor, + ) -> None: + # Broadcasting accepts a bias of the wrong length-1 shape and a gate of fewer + # dims; mixed dtypes promote and change every rounding. None of these raise. + assert attention_output.dim() == 3 and gate.shape == attention_output.shape, ( + "attention_output and gate must both be [T, heads, head_dim]" + ) + assert gate_bias.shape == (attention_output.shape[1],), ( + f"gate_bias must be [heads] = [{attention_output.shape[1]}], " + f"got {tuple(gate_bias.shape)}" + ) + assert attention_output.dtype == gate.dtype == gate_bias.dtype, ( + "attention_output, gate and gate_bias must share one dtype" + ) + + +beam_output_gate = BeamOutputGate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py index 80a2ad211f21..72edf3f716c7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py @@ -92,6 +92,55 @@ def _spread(batch: int, longest: int) -> tuple[int, ...]: return tuple(1 + (longest - 1) * i // (batch - 1) for i in range(batch)) +#: The Beam dense block's sliding-attention layer: GQA 48q/8kv d128 over a bf16 pool, +#: RoPE applied outside, one fresh context per call on the packed path. The window is +#: 512 and so is the pool's per-sequence capacity (16 pages of 32), so the window equals +#: max_seq_len and no context the block accepts reaches it. +_BEAM_DENSE = dict( + input_type=STANDARD, + num_heads=48, + num_kv_heads=8, + head_size=128, + v_head_dim=128, + sinks=False, + window=512, + quant_mode=0, + mask_type=1, + tokens_per_block=32, + num_blocks=16, + max_num_requests=1, + use_paged_context_fmha=False, + num_contexts=1, + predicted_tokens_per_seq=1, +) + +#: The Beam six-block model's layers as its generator runs them: the same geometry over a +#: cache that already holds the request's earlier tokens, on the paged-context path (block +#: reuse is on), with a 640-token per-sequence pool (20 pages) so histories cross the +#: 512-token window. `_BEAM_FULL` is block 5, the full-attention layer. +_BEAM_SLIDING = dict( + _BEAM_DENSE, + num_blocks=64, + max_num_requests=3, + max_blocks_per_seq=20, + use_paged_context_fmha=True, +) +_BEAM_FULL = dict(_BEAM_SLIDING, window=None) + +#: The Beam full model's layers on one of four TP ranks (12q/2kv d128) as a DFlash2 +#: verification step runs them: P = 8 rows per generation sequence (the anchor and seven +#: proposals), on the paged-context path, with a 768-token per-sequence pool (24 pages) so +#: histories cross the 512-token window. `_BEAM_TP4_FULL` is a full-attention layer. +_BEAM_TP4_SLIDING = dict( + _BEAM_SLIDING, + num_heads=12, + num_kv_heads=2, + num_blocks=96, + max_num_requests=4, + max_blocks_per_seq=24, +) +_BEAM_TP4_FULL = dict(_BEAM_TP4_SLIDING, window=None) + #: A GQA 24/4/256 full-attention layer (Qwen3.8-27B) over an fp8-e4m3 pool at #: the production kv scale 1.0. Same chosen pool sizes as `_QWEN4_DENSE`. _FP8_24_4_256 = dict( @@ -765,6 +814,101 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=4, ), ), + *( + Cell( + why=f"beam dense block: GQA 48q/8kv d128, fresh {tokens}-row context", + spec=dict(_BEAM_DENSE, cached=(0,), new=(tokens,)), + ) + for tokens in (1, 17, 129) + ), + Cell( + why="beam sliding, paged context: 17 rows over a 64-token reused prefix", + spec=dict( + _BEAM_SLIDING, cached=(64,), new=(17,), num_contexts=1, predicted_tokens_per_seq=1 + ), + ), + Cell( + why="beam sliding, paged context: a fresh 513-row context, key 0 leaves the " + "window at row 512", + spec=dict( + _BEAM_SLIDING, cached=(0,), new=(513,), num_contexts=1, predicted_tokens_per_seq=1 + ), + ), + Cell( + why="beam sliding, paged context: one row at position 512 after a 512-token " + "reused prefix", + spec=dict( + _BEAM_SLIDING, cached=(512,), new=(1,), num_contexts=1, predicted_tokens_per_seq=1 + ), + ), + Cell( + why="beam sliding decode across the window: histories 511, 512 and 513", + spec=dict( + _BEAM_SLIDING, + cached=(511, 512, 513), + new=(1, 1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="beam full layer decode past the sliding window: histories 511, 512 and 513", + spec=dict( + _BEAM_FULL, + cached=(511, 512, 513), + new=(1, 1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="beam full layer, paged context: 17 rows over a 64-token reused prefix", + spec=dict( + _BEAM_FULL, cached=(64,), new=(17,), num_contexts=1, predicted_tokens_per_seq=1 + ), + ), + Cell( + why="beam tp4 sliding verify P=8: blocks after histories 503, 504 and 505 end " + "at 511, 512 and 513, crossing the 512-token window", + spec=dict( + _BEAM_TP4_SLIDING, + cached=(503, 504, 505), + new=(8, 8, 8), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="beam tp4 full verify P=8: the same blocks on a full-attention layer", + spec=dict( + _BEAM_TP4_FULL, + cached=(503, 504, 505), + new=(8, 8, 8), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="beam tp4 sliding mixed step P=8: a 20-row context over a 31-token prefix " + "ahead of three blocks over histories 17, 600 and 129", + spec=dict( + _BEAM_TP4_SLIDING, + cached=(31, 17, 600, 129), + new=(20, 8, 8, 8), + num_contexts=1, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="beam tp4 full mixed step P=8: the same step on a full-attention layer", + spec=dict( + _BEAM_TP4_FULL, + cached=(31, 17, 600, 129), + new=(20, 8, 8, 8), + num_contexts=1, + predicted_tokens_per_seq=8, + ), + ), ) note = """ @@ -785,6 +929,11 @@ class its own layer group, so its mapping carries two pool ids; the op reads beside this one certify, and what the matrix here checks by reading the pool back rather than by a reference. + A sliding window is a mask, not a ring buffer: the append writes every token at its + absolute position whatever `attention_window_size` is, and a query reads the keys of its + window from those positions (the Beam window cells read the pool back past position 512). + Pages wholly below every query's window are never read. + The softmax statistics the MLA chunked-prefill path writes are a second output; `softmax_stats` states them separately so `reference` stays one thing. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py index eccf450ef885..50d8c3e7e1d9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py @@ -10,6 +10,32 @@ from .._op import Arch, Cell, OpWrapper, true_fp32_matmul +#: The Beam dense block's bias-free projections at its three fresh-context lengths: q and +#: the attention gate 5120 -> 6144, k and v 5120 -> 1024, o 6144 -> 5120, the feed-forward +#: gate and up 5120 -> 20480, and down 20480 -> 5120. +_BEAM_DENSE_BLOCK = tuple( + Cell( + why=f"beam dense {name} {k}->{n}, {m} tokens", + spec=dict(m=m, k=k, n=n, dtype=torch.bfloat16, bias=False, out_dtype=None), + ) + for name, k, n in ( + ("q/gate", 5120, 6144), + ("k/v", 5120, 1024), + ("o", 6144, 5120), + ("gate/up", 5120, 20480), + ("down", 20480, 5120), + ) + for m in (1, 17, 129) +) +#: The Beam six-block model's bias-free, untied head: 5120 -> the 151936-entry vocabulary. +_BEAM_HEAD = tuple( + Cell( + why=f"beam head 5120->151936, {m} tokens", + spec=dict(m=m, k=5120, n=151936, dtype=torch.bfloat16, bias=False, out_dtype=None), + ) + for m in (1, 17, 129) +) + class CublasMm(OpWrapper): """`mat_a @ mat_b (+ bias)`, fp32-accumulated, one cublas_mm call. @@ -24,7 +50,8 @@ class CublasMm(OpWrapper): # three 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). + # asks for fp32 logits). The Beam dense block's cells are by size: every + # projection it makes, at each of its context lengths. CELLS: tuple[Cell, ...] = ( Cell( why="gpt-oss qkv with fused bias: hidden 2880 -> (64 + 8 + 8) * 64", @@ -44,6 +71,8 @@ class CublasMm(OpWrapper): why="decode-shaped: one token, the shape every projection sees per step", spec=dict(m=1, k=7168, n=1536, dtype=torch.bfloat16, bias=False, out_dtype=None), ), + *_BEAM_DENSE_BLOCK, + *_BEAM_HEAD, Cell( why="deepseek default q_b at attention TP4: q_lora 1536 -> 32 heads x 192", spec=dict(m=16, k=1536, n=6144, dtype=torch.bfloat16, bias=False, out_dtype=None), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py index b6280682aa7a..3629373627ea 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py @@ -16,6 +16,22 @@ from .._op import Arch, Cell, OpWrapper +#: The Beam dense block's norms at its three fresh-context lengths: the four 5120-wide +#: sandwich norms, and the per-head q and k norms, which normalize every 128-wide head of +#: a token as its own row -- 48 q rows and 8 k rows per token. +_BEAM_DENSE_BLOCK = tuple( + Cell( + why=f"beam dense {name}, {tokens} tokens: {rows} rows of {width}", + spec=dict(shape=(rows, width), dtype=torch.bfloat16, eps=1e-6), + ) + for tokens in (1, 17, 129) + for name, rows, width in ( + ("sandwich norm", tokens, 5120), + ("q_norm", tokens * 48, 128), + ("k_norm", tokens * 8, 128), + ) +) + class FlashinferRmsnorm(OpWrapper): """`x / sqrt(mean(x^2, -1) + eps) * weight`, one op call, new tensor out. @@ -31,7 +47,8 @@ class FlashinferRmsnorm(OpWrapper): # only on its residual stream; deepseek reaches it there and on the two # lora ranks its MLA projects through, which is why the narrow cells are # here at all -- a 512-wide row is one 128-bit vector per lane and takes a - # different path through the kernel than a 7168-wide one. + # different path through the kernel than a 7168-wide one. The Beam dense + # block adds 5120-wide rows and, for its per-head q/k norms, 128-wide ones. CELLS: tuple[Cell, ...] = ( Cell( why="decode-shaped, gpt-oss-120b hidden 2880", @@ -53,6 +70,7 @@ class FlashinferRmsnorm(OpWrapper): why="deepseek kv_lora rank 512, the narrowest row any target asks for", spec=dict(shape=(16, 512), dtype=torch.bfloat16, eps=1e-6), ), + *_BEAM_DENSE_BLOCK, ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/__init__.py new file mode 100644 index 000000000000..51c47b7ad054 --- /dev/null +++ b/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.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/beam_depth_residual.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/beam_depth_residual.py new file mode 100644 index 000000000000..4706d71f2de8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/beam_depth_residual.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Depth-scaled residual add, in the activation dtype.""" + +from __future__ import annotations + +import math + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: A block's two adds as (branch, residual) scales: (1, sqrt(n - 1)) / sqrt(n) for the n-th +#: add of the residual stream. Block i's attention add is n = 2(i + 1) and its feed-forward +#: add the next one: n = 2 and 3 for block 0, n = 8 and 9 for block 3. +_ATTENTION_ADD = (1.0 / math.sqrt(2.0), math.sqrt(1.0) / math.sqrt(2.0)) +_FEED_FORWARD_ADD = (1.0 / math.sqrt(3.0), math.sqrt(2.0) / math.sqrt(3.0)) +_BLOCK3_ATTENTION_ADD = (1.0 / math.sqrt(8.0), math.sqrt(7.0) / math.sqrt(8.0)) +_BLOCK3_FEED_FORWARD_ADD = (1.0 / math.sqrt(9.0), math.sqrt(8.0) / math.sqrt(9.0)) + +_HIDDEN = dict(hidden=5120, dtype=torch.bfloat16) + + +class BeamDepthResidual(OpWrapper): + """`branch * branch_scale + residual * residual_scale`, a new tensor. + + Three activation-dtype operations -- two scalings and the add -- each rounded to + that dtype, as the source model computes them. The scales are Python floats and + reach the kernel as fp32. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="block 0 attention add, 17 tokens", + spec=dict(_HIDDEN, tokens=17, scales=_ATTENTION_ADD), + ), + Cell( + why="block 0 feed-forward add, 1 token", + spec=dict(_HIDDEN, tokens=1, scales=_FEED_FORWARD_ADD), + ), + Cell( + why="block 0 feed-forward add, 129 tokens: the longest context", + spec=dict(_HIDDEN, tokens=129, scales=_FEED_FORWARD_ADD), + ), + Cell( + why="block 3 attention add, 17 tokens", + spec=dict(_HIDDEN, tokens=17, scales=_BLOCK3_ATTENTION_ADD), + ), + Cell( + why="block 3 feed-forward add, 17 tokens", + spec=dict(_HIDDEN, tokens=17, scales=_BLOCK3_FEED_FORWARD_ADD), + ), + ) + + note = """ + Not fused with the norm before it: the source rounds the norm output, then each + scaled term, then the sum. A fused fp32 add-and-scale rounds once and is a + different result. + """ + + def __call__( + self, + branch: torch.Tensor, + residual: torch.Tensor, + branch_scale: float, + residual_scale: float, + ) -> torch.Tensor: + return branch * branch_scale + residual * residual_scale + + def reference( + self, + branch: torch.Tensor, + residual: torch.Tensor, + branch_scale: float, + residual_scale: float, + ) -> torch.Tensor: + """Each product and the sum in fp32, rounded to the activation dtype in turn.""" + dtype = branch.dtype + scaled_branch = (branch.float() * branch_scale).to(dtype) + scaled_residual = (residual.float() * residual_scale).to(dtype) + return (scaled_branch.float() + scaled_residual.float()).to(dtype) + + def is_valid( + self, + branch: torch.Tensor, + residual: torch.Tensor, + branch_scale: float, + residual_scale: float, + ) -> None: + # Broadcasting accepts mismatched shapes, mixed dtypes promote, and a tensor + # scale in the activation dtype rounds the coefficient itself. None raise. + assert branch.shape == residual.shape, ( + f"branch {tuple(branch.shape)} and residual {tuple(residual.shape)} must match" + ) + assert branch.dtype == residual.dtype, "branch and residual must share one dtype" + assert isinstance(branch_scale, float) and isinstance(residual_scale, float), ( + "scales must be Python floats" + ) + + +beam_depth_residual = BeamDepthResidual() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md new file mode 100644 index 000000000000..d1111389e62f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md @@ -0,0 +1,159 @@ +# Beam with ModelingV2 + +This experimental target runs the complete 52-layer Beam model through TRT-LLM's +native PyTorch executor using the StairCase/ModelingV2 approach. It includes dense +and MoE decoder blocks, checkpoint loading, cached generation, shared-prefix KV +reuse, and a text chat-completions endpoint. The supported launch below requires +ModelingV2 explicitly; it does not fall back to the built-in model implementation. + +## Supported deployment + +| Setting | Supported configuration | +|---|---| +| Hardware | Four GB300 GPUs (SM103) on one node | +| Parallel layout | Attention TP4, MoE TP4/EP1, PP1; attention DP disabled | +| Parameters and arithmetic | Packed NVFP4 experts with native W4A4 GEMMs; BF16 excluded modules/activations; FP32 routing | +| Attention | TRTLLM backend; 512-token sliding attention and full attention according to the checkpoint's layer schedule | +| KV cache | BF16, native `KVCacheManagerV2`, full-block prefix reuse | +| Generation | Native `TorchSampler` for ordinary decoding; native DFlash worker and speculative sampler for optional DFlash2 greedy decoding | +| Execution | Eager; CUDA graphs, overlap scheduling, and chunked prefill disabled | +| Serving | Native OpenAI-compatible server, one frontend, streaming and non-streaming text | + +The expert provider currently dispatches selected experts in a host-driven loop. +The GEMMs use native NVFP4 kernels, but expert dispatch and rank agreement introduce +synchronization. This is a functional baseline, not a throughput-optimized MoE path. + +## Launch and request + +Use a TensorRT-LLM installation built from the commit containing this target, with +matching native libraries and access to all four supported GPUs. Supply an +authorized, compatible checkpoint directory containing its original configuration, +tokenizer, chat template, and indexed safetensors weights. No checkpoint download +or conversion is performed by this recipe. + +```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 launcher selects the target, checkpoint loader, tokenizer, channel parser, +cache manager, and execution settings together. Plain `trtllm-serve` with default +settings is not an equivalent launch recipe. Analysis-channel text is returned as +`reasoning_content`; final-channel text is returned as `content`. + +```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}}' +``` + +Set `stream` to `true` for streaming. `max_tokens` includes reasoning tokens, so a +short budget can terminate before final-channel text is produced. The request above +is an invocation example, not a promise of a particular response or token length. + +The sequence limit includes prompt and generated tokens. With chunked prefill off, +the prompt must also fit the per-step token budget; the model accepts at most 8192 +new token rows per call. Configured sequence capacity does not establish accuracy +at the maximum sequence length. The launcher defaults to one request per step; +native batching can be enabled with `--max-batch-size`. The benchmark used 128 and +`--kv-cache-max-tokens 1048576`, with workload-specific memory budgeting. Shared KV +capacity must cover concurrently active requests; do not infer capacity from batch +size alone. `--iter-stats` enables native iteration statistics. + +## Optional DFlash2 speculative decoding + +Supply a compatible DFlash2 drafter directory with `--draft-checkpoint`. The +launcher uses the native DFlash worker, greedy verifier, and separate target/draft +KV cache pools. The target supplies completed-block captures at the layers named +by the drafter; the drafter shares the target's embedding and vocabulary head. + +```bash +python -m tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.serving \ + --checkpoint /path/to/target-checkpoint \ + --draft-checkpoint /path/to/dflash2-checkpoint \ + --served-model-name beam-nvfp4 \ + --host 127.0.0.1 --port 8000 \ + --max-seq-len 8192 --max-num-tokens 8192 --max-batch-size 4 +``` + +The evaluated configuration uses four GB300 GPUs, TP4, BF16 KV for both models, +seven proposals per step (eight target verification rows), and eager execution +without overlap or chunked prefill. Validation covers batches of one and four, +prompts through 4097 tokens, and at most 64 generated tokens; the configured 8192 +sequence capacity is not a claim of qualification at every length. The default +shared KV capacity can queue a request when four long prompts compete for space. + +Draft context uses the native paged cache. Non-causal windowed draft block +attention uses Torch SDPA over that same pool because TRTLLM-Gen does not support +that mask. Target verification retains TRTLLM attention. Accepted proposals must +match the target's greedy choice on the verification logits; independent ordinary +and speculative runs can still diverge at near ties when GEMM/collective shapes +change. The GSM8K result below is for ordinary decoding, not DFlash2. + +FP8 KV, CUDA graphs, overlap scheduling, and disaggregated serving remain future +work. The bounded eager timing measurements are not a production throughput claim. + +## Validation and regression coverage + +The initial full-model evaluation completed all 1,319 GSM8K test questions with +10-shot prompting, greedy decoding, reasoning effort 128, and a 65,536-token +completion allowance. It scored **1,263/1,319 (95.75%)**, with no execution errors +or truncated completions. Evaluation used disjoint subsets on two four-GPU +replicas, followed by resumptions of unfinished questions; completed samples were +retained and each question was scored once. + +The comparison baseline is the recorded BF16 mean of **95.53%** across three runs, +not a same-precision NVFP4 reference. Precision, runtime, and batching differ. +Greedy token trajectories are not guaranteed to be identical across batch sizes, +and this small score difference is not evidence that quantization improves +accuracy. Private checkpoint identifiers, evaluation traces, and reference captures +are not part of this source distribution. + +Portable generated-input tests cover the component implementations and their +composition. The full-model directory contains CPU parser/routing tests and GPU +metadata/request-row tests: + +```bash +python -m pytest -m cpu_only \ + tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_serving.py +python -m pytest \ + tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_batching.py +``` + +The second command requires a compatible CUDA environment but no checkpoint. +These unit tests do not reproduce the full 52-layer accuracy evaluation. A +checkpoint-backed end-to-end regression and its CI model staging remain separate +integration work; generated tests must not be presented as that coverage. + +L0 includes the CPU parser/routing and attention-selection checks, plus the +single-GPU generated-input operator, block, six-layer and metadata regressions. +The GPU tests are selected by the ModelingV2 directory entry in `l0_b300.yml`; +the six-layer suites need up to approximately 200 GiB of free device memory and +release their cached models at module teardown. Real-checkpoint multi-GPU smoke +coverage is outside L0 and may run manually or in a separately provisioned +nightly/long-running job. Full Beam GSM8K evaluation is manual-only and must not +be registered in L0, nightly or long-running CI. No Beam multi-GPU CI job is +registered by this change. + +## File organization and limits + +- `full_model_nvfp4__sm_103__tp4/` contains the engine-facing target, full-depth + execution, native checkpoint loader adapter, and serving launcher/parser. +- `checkpoint_layout_nvfp4__sm_103__tp4/` supplies packed weights and expert kernels; + it reuses BF16 layout/loading and block helpers from + `checkpoint_layout_bf16__sm_103__tp4/`. +- Router, expert, dispatch/combine, dense/sparse block, and six-layer directories + retain component implementations and generated-input regression coverage. Some + are runtime dependencies of the full target, not disposable tier artifacts. +- `tensorrt_llm/_torch/models/modeling_beam.py` supplies `BeamConfig` and the earlier + built-in component implementation. The result above belongs to the routed + ModelingV2 target, not to that built-in model path. + +Other GPU architectures, parallel layouts, FP8 KV cache, other speculative modes, +CUDA graphs, overlap scheduling, chunked prefill, tool calling, and disaggregated +serving are not qualified by this implementation. Subsequent feature/performance +work should preserve this functional and accuracy baseline in follow-up PRs. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/__init__.py new file mode 100644 index 000000000000..890bdb17925a --- /dev/null +++ b/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.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py new file mode 100644 index 000000000000..ba71826dec2c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py @@ -0,0 +1,493 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The attention sublayer of the Beam decoder block components, one request step. + +Shared by the dense and sparse block components. It holds no parameters, cache or +configuration: each block passes its own tensors and geometry, its layer's position +encoding and window, and the prepared attention metadata of the call. The computation, +before the block's post-attention norm:: + + q, k, v, g = h @ q_proj^T, h @ k_proj^T, h @ v_proj^T, h @ gate_proj^T + q, k = rmsnorm(q, q_norm), rmsnorm(k, k_norm) per head, rounded to bf16 + sliding layer: q, k = rope(q, k) fp32 NeoX rotation, rounded + full layer: q = q * exp(0.1 * log1p((p + 1) / 2048)) fp32 scale, rounded; k as is + a = causal attention(q, k, v), scale 1/sqrt(128), 6 q heads per kv head, + window 512 on a sliding layer, the whole context on a full one + a = a * sigmoid(g + gate_bias) per head, bf16 steps + out = a @ o_proj^T + +A full-attention layer has no position embedding: neither q nor k is rotated, and its +only position dependence is the per-token query scale above. + +A step of one request: a context call (``num_contexts == 1``) appends its T new tokens +after the ``past`` tokens the cache already holds for the request -- none for a fresh +context, a reused or earlier-computed prefix otherwise -- and a generation call +(``num_contexts == 0``) appends exactly one token. Every position is absolute: the new +tokens sit at ``past .. past + T - 1``, where ``past + T`` is the request's KV length in +the metadata, and the positions the block receives must be exactly those. A context call +over a cached prefix must use paged-context FMHA, the path that reads the prefix's keys and +values from the pool; without it the context kernel would attend to the new rows only. +The keys and values the op appends are the transformed k (normalized, rotated on a sliding +layer) and the projected v of the new tokens, written to the request's own pages. + +A step of several requests (``check_batch_metadata``) is the native runtime's batch: its +context requests first, then its generation requests, each with its own ``past``, KV length +and pages, their new tokens concatenated in that order. Every operation but the attention +op is per token; the op reads each request's own metadata row. + +A speculative verification step (one-model speculative decoding) is the same batch whose +generation requests each append ``P`` rows instead of one: the pending token and its +``P - 1`` proposals at ``past .. past + P - 1``. The op takes ``P`` as +``predicted_tokens_per_seq`` (``call_constants``), which alone makes the mask inside each +block bottom-right causal: row ``t`` attends ``[0, past + t]`` (within the layer's window). +On sm_103 the metadata keeps ``is_spec_decoding_enabled`` False for such a linear draft. The +one-model runtime reserves extra KV slots per request (``num_extra_kv_tokens``) beyond the +lengths the op reads; the KV lengths in the metadata exclude them. + +Logical layer and cache slot. A block's ``layer_idx`` is its depth in the model: it fixes +the block's semantics (its residual coefficients, its attention and feed-forward kind). +The KV cache it appends to is a ``KVCacheManagerV2`` that holds that layer, so its storage +slot is the manager's local index for it -- ``layer_offsets[layer_idx]``: 0 for the one +layer a ``layer_mask`` selects, the layer's own position in a manager holding several. +The attention op and the manager's per-layer arrays take that local index; manager APIs +such as ``get_buffers`` take the logical one and translate it. The slot is resolved from +each call's metadata, never stored, and a manager that does not hold the block's logical +layer is rejected before any kernel runs. + +Host synchronization: the one-request position check (``check_step``) compares the +positions with ``arange(past, past + T)`` on the device and reads the answer back, once per +call. The metadata checks read host tensors only. +""" + +from __future__ import annotations + +from typing import Mapping + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_neox_rope import ( + beam_neox_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_nope_query_scale import ( # noqa: E501 + beam_nope_query_scale, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_output_gate import ( + beam_output_gate, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + STANDARD, + thop_attention, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm.bindings import DataType + +#: thop_attention arguments that do not change with the step on these blocks: packed q|k|v, +#: append-and-attend, causal, bf16 pool, no quantization, and any position transform +#: applied before the call -- the rope_* values are inert placeholders the op requires, +#: not this model's RoPE. +CALL_CONSTANTS = dict( + k=None, + v=None, + output_sf=None, + is_fused_qkv=True, + update_kv_cache=True, + predicted_tokens_per_seq=1, + attention_input_type=STANDARD, + is_mla_enable=False, + mask_type=1, + q_scaling=1.0, + quant_mode=0, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + attention_sinks=None, + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + out_scale=None, + latent_cache=None, + q_pe=None, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + chunked_prefill_buffer_batch_size=1, + attention_chunk_size=None, + softmax_stats_tensor=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + cross_kv=None, + relative_attention_bias=None, + relative_attention_max_distance=0, + kv_norm_weight=None, + kv_norm_eps=1e-6, + skip_correction_threshold=0.0, +) + + +def call_constants(generation_rows: int = 1) -> dict: + """``CALL_CONSTANTS`` for a step whose generation requests each append ``generation_rows``. + + 1 is an ordinary step; a speculative verification step passes its ``P``. + """ + if type(generation_rows) is not int or generation_rows < 1: + raise ValueError(f"generation_rows must be a positive int, got {generation_rows!r}") + return {**CALL_CONSTANTS, "predicted_tokens_per_seq": generation_rows} + + +def local_layer_index(md: TrtllmAttentionMetadata, layer_idx: int) -> int: + """The cache slot of logical layer ``layer_idx`` in the call's cache manager. + + Raises ValueError unless the manager holds ``layer_idx`` at an integer local index + within its local layers. Reads the manager's host-side layer map only. + """ + mgr = md.kv_cache_manager + offsets = getattr(mgr, "layer_offsets", None) + if not isinstance(offsets, Mapping) or layer_idx not in offsets: + held = sorted(offsets) if isinstance(offsets, Mapping) else None + raise ValueError( + f"the KV cache manager does not hold logical layer {layer_idx}; it holds {held}. " + f"Select the layer with a layer_mask rather than reusing another layer's cache" + ) + num_local = getattr(mgr, "num_local_layers", None) + local = offsets[layer_idx] + if ( + type(num_local) is not int + or type(local) is not int + or len(offsets) != num_local + or not 0 <= local < num_local + ): + raise ValueError( + f"the KV cache manager maps logical layer {layer_idx} to slot {local!r} of " + f"{num_local!r} local layers holding {len(offsets)} logical layers" + ) + return local + + +def check_step( + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + md: TrtllmAttentionMetadata, + *, + layer_idx: int, + hidden_size: int, + num_kv_heads: int, + head_dim: int, + device: torch.device, + max_tokens: int, +) -> None: + """Raise ValueError on any call input a block would otherwise compute wrongly. + + Accepted: ``hidden_states`` contiguous ``[1, T, hidden_size]`` bf16 on ``device``, + 1 <= T <= ``max_tokens``; integer ``position_ids`` ``[1, T]`` equal to + ``arange(past, past + T)`` on that device; a prepared ``TrtllmAttentionMetadata`` + holding one step of one request on a BF16 ``KVCacheManagerV2`` that holds logical + layer ``layer_idx`` with ``num_kv_heads`` kv heads of ``head_dim`` (``check_metadata`` + states the steps and returns ``past``). + """ + if not isinstance(hidden_states, torch.Tensor) or hidden_states.dtype != torch.bfloat16: + raise ValueError("hidden_states must be a torch.bfloat16 tensor") + if ( + hidden_states.dim() != 3 + or hidden_states.shape[0] != 1 + or hidden_states.shape[2] != hidden_size + ): + raise ValueError( + f"hidden_states must be [1, T, {hidden_size}], got {tuple(hidden_states.shape)}" + ) + if hidden_states.device != device or not hidden_states.is_contiguous(): + raise ValueError("hidden_states must be contiguous on the block's device") + num_tokens = hidden_states.shape[1] + if not 1 <= num_tokens <= max_tokens: + raise ValueError(f"supported token counts per call are 1..{max_tokens}, got {num_tokens}") + if ( + not isinstance(position_ids, torch.Tensor) + or position_ids.is_floating_point() + or tuple(position_ids.shape) != (1, num_tokens) + or position_ids.device != hidden_states.device + ): + raise ValueError( + f"position_ids must be an integer [1, {num_tokens}] tensor on the block's device" + ) + past = check_metadata( + md, num_tokens, layer_idx=layer_idx, num_kv_heads=num_kv_heads, head_dim=head_dim + ) + expected = torch.arange( + past, past + num_tokens, dtype=position_ids.dtype, device=position_ids.device + ) + if not torch.equal(position_ids[0], expected): + raise ValueError( + f"the call appends tokens {past}..{past + num_tokens - 1} after {past} cached " + "tokens; its positions must be exactly those" + ) + + +def check_metadata( + md: TrtllmAttentionMetadata, + num_tokens: int, + *, + layer_idx: int, + num_kv_heads: int, + head_dim: int, +) -> int: + """Raise ValueError unless ``md`` is one prepared step of one request; return ``past``. + + ``check_batch_metadata`` restricted to a batch of exactly one request: a context step of + ``num_tokens`` new tokens after ``past`` cached ones, or a generation step of one token. + ``past`` is the absolute position of the step's first new token. Reads host-side state + only. + """ + spans = check_batch_metadata( + md, num_tokens, layer_idx=layer_idx, num_kv_heads=num_kv_heads, head_dim=head_dim + ) + if len(spans) != 1: + raise ValueError(f"one request per call is supported; metadata has {len(spans)}") + return spans[0][0] + + +def check_batch_metadata( + md: TrtllmAttentionMetadata, + num_tokens: int, + *, + layer_idx: int, + num_kv_heads: int, + head_dim: int, + generation_rows: int = 1, + extra_kv_tokens: int = 0, +) -> list: + """Raise ValueError unless ``md`` is one prepared step of its requests; return their spans. + + The native runtime orders a step's ``num_seqs`` requests with its ``num_contexts`` + context requests first, and their new tokens fill the call's ``num_tokens`` rows in that + order. Accepted on a BF16 ``KVCacheManagerV2`` holding logical layer ``layer_idx`` with + ``num_kv_heads`` kv heads of ``head_dim``, beam width 1, no cross attention and no + tree-mask speculative decoding, where each request ``i`` with KV length ``L_i`` (at most + the manager's ``max_seq_len``) is + + * a context request of ``n_i >= 1`` new tokens: context length ``n_i``, + ``past_i = L_i - n_i >= 0`` tokens already cached, and paged-context FMHA whenever + ``past_i > 0``; or + * a generation request of ``generation_rows`` new tokens after ``past_i = L_i - n_i >= 0`` + cached ones (the native runtime's warmup issues generation steps over a cache that + holds only the new tokens): one for an ordinary step, ``P`` for a speculative + verification step (the pending token and its ``P - 1`` proposals). + + ``extra_kv_tokens`` is the per-request KV reservation the metadata must carry beyond the + lengths above: 0 for an ordinary engine, the one-model runtime's ``num_extra_kv_tokens`` + for a speculative one. + + Returns ``[(past_i, n_i), ...]`` in row order: request ``i``'s new tokens sit at the + absolute positions ``past_i .. past_i + n_i - 1``. Reads host-side state only. + """ + if not isinstance(md, TrtllmAttentionMetadata): + raise ValueError(f"attn_metadata must be TrtllmAttentionMetadata, got {type(md).__name__}") + if type(generation_rows) is not int or generation_rows < 1: + raise ValueError(f"generation_rows must be a positive int, got {generation_rows!r}") + if type(extra_kv_tokens) is not int or extra_kv_tokens < 0: + raise ValueError(f"extra_kv_tokens must be a non-negative int, got {extra_kv_tokens!r}") + mgr = md.kv_cache_manager + if mgr is None: + raise ValueError("attn_metadata needs a KV cache manager for the paged append") + if md.is_cross or md.is_spec_decoding_enabled or md.beam_width != 1: + raise ValueError( + "cross attention, tree-mask speculative decoding and beams are not supported" + ) + num_seqs, num_contexts = md.num_seqs, md.num_contexts + if num_seqs < 1 or not 0 <= num_contexts <= num_seqs: + raise ValueError( + f"the metadata holds {num_seqs} requests of which {num_contexts} are contexts" + ) + if md.kv_cache_params is None: + raise ValueError("the metadata must use the KV cache") + reserved = md.kv_cache_params.num_extra_kv_tokens or 0 + if reserved != extra_kv_tokens: + raise ValueError( + f"the metadata reserves {reserved} extra KV tokens per request; this step needs " + f"{extra_kv_tokens} (0 without speculative decoding)" + ) + if num_tokens > md.max_num_tokens: + raise ValueError( + f"the call has {num_tokens} tokens; the metadata holds {md.max_num_tokens} per call" + ) + new_tokens = [int(n) for n in md.seq_lens[:num_seqs]] + kv_lens = [int(n) for n in md.kv_lens_runtime[:num_seqs]] + request_types = [int(t) for t in md.host_request_types_runtime[:num_seqs]] + context_lens = [int(n) for n in md.prompt_lens_cpu_runtime[:num_contexts]] + if sum(new_tokens) != num_tokens or md.num_ctx_tokens != sum(new_tokens[:num_contexts]): + raise ValueError( + f"the requests' new tokens {new_tokens} ({md.num_ctx_tokens} in contexts) do not " + f"fill the call's {num_tokens} tokens" + ) + spans = [] + for i, (n, kv_len, request_type) in enumerate(zip(new_tokens, kv_lens, request_types)): + past = kv_len - n + if i < num_contexts: + if n < 1 or context_lens[i] != n or request_type != 0: + raise ValueError( + f"context request {i} of {n} new tokens needs context length {n}; " + f"metadata has context length {context_lens[i]} and request type " + f"{request_type}" + ) + if past < 0: + raise ValueError( + f"request {i}'s KV length {kv_len} is shorter than its {n} new tokens" + ) + if past and not md.use_paged_context_fmha: + raise ValueError( + f"a context step after {past} cached tokens must read them from the pool: " + "paged-context FMHA is required" + ) + else: + if n != generation_rows or request_type != 1: + rows = "one token" if generation_rows == 1 else f"{generation_rows} tokens" + raise ValueError( + f"generation request {i} appends exactly {rows} this step, got {n} " + f"(request type {request_type})" + ) + if past < 0: + raise ValueError( + f"request {i}'s KV length {kv_len} is shorter than its {n} generation rows" + ) + spans.append((past, n)) + if max(kv_lens) > mgr.max_seq_len or md.max_seq_len != mgr.max_seq_len: + raise ValueError( + f"a request reaches {max(kv_lens)} tokens; the cache holds {mgr.max_seq_len} per " + f"sequence and the metadata states {md.max_seq_len}" + ) + local = local_layer_index(md, layer_idx) + heads = getattr(mgr, "num_kv_heads_per_layer", None) or [] + local_heads = heads[local] if len(heads) == mgr.num_local_layers else None + if mgr.dtype != DataType.BF16 or local_heads != num_kv_heads or mgr.head_dim != head_dim: + raise ValueError( + f"the KV cache must be BF16 with {num_kv_heads} kv heads of {head_dim}; " + f"got {mgr.dtype}, {heads}, head_dim {mgr.head_dim}" + ) + return spans + + +def step_args(md: TrtllmAttentionMetadata, layer_idx: int) -> dict: + """Project the prepared metadata onto thop_attention's explicit batch state. + + Includes ``local_layer_idx``, the cache slot of logical layer ``layer_idx`` resolved + from this call's manager. + """ + return dict( + local_layer_idx=local_layer_index(md, layer_idx), + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_total_kv_lens=md.host_total_kv_lens, + context_lengths=md.prompt_lens_cuda_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + host_request_types=md.host_request_types_runtime, + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, + workspace_=md.effective_workspace, + tokens_per_block=md.tokens_per_block, + max_num_requests=md.max_num_requests, + max_context_length=md.max_context_length, + max_seq_len=md.max_seq_len, + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, + use_paged_context_fmha=md.use_paged_context_fmha, + beam_width=md.effective_beam_width, + cache_indirection=md.cache_indirection, + block_ids_per_seq=md.block_ids_per_seq, + max_context_q_len_override=md.max_context_q_len_override, + is_cross=md.is_cross, + is_spec_decoding_enabled=md.is_spec_decoding_enabled, + use_spec_decoding=md.use_spec_decoding, + is_spec_dec_tree=md.is_spec_dec_tree, + spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, + spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp, + spec_decoding_packed_mask=md.spec_decoding_packed_mask, + spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, + spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, + spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + num_sparse_topk=md.num_sparse_topk, + flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, + flash_mla_num_splits=md.flash_mla_num_splits, + max_num_sequences=md.max_num_sequences, + force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + ) + + +def attention( + h: torch.Tensor, + positions: torch.Tensor, + step: dict, + *, + q_proj: torch.Tensor, + k_proj: torch.Tensor, + v_proj: torch.Tensor, + gate_proj: torch.Tensor, + o_proj: torch.Tensor, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + gate_bias: torch.Tensor, + num_heads: int, + num_kv_heads: int, + head_dim: int, + eps: float, + rope_theta: float | None, + window: int, +) -> torch.Tensor: + """Normalized input [T, hidden] -> attention output [T, hidden], before post_attention_layernorm. + + ``step`` is ``step_args`` of the call's metadata; the projections are ``[out, in]``. + ``rope_theta`` is the sliding layer's RoPE base; ``None`` is a full-attention layer + without position embedding, whose q takes the per-token query scale instead. ``window`` + is the attention window the op applies: the sliding window, or the metadata's + ``max_seq_len`` for a full layer. + """ + num_tokens = h.shape[0] + heads, kv_heads, d = num_heads, num_kv_heads, head_dim + q = cublas_mm(h, q_proj.t()) + k = cublas_mm(h, k_proj.t()) + v = cublas_mm(h, v_proj.t()) + gate = cublas_mm(h, gate_proj.t()) + q = flashinfer_rmsnorm(q.view(num_tokens * heads, d), q_norm, eps) + k = flashinfer_rmsnorm(k.view(num_tokens * kv_heads, d), k_norm, eps) + q, k = q.view(num_tokens, heads, d), k.view(num_tokens, kv_heads, d) + if rope_theta is None: + q = beam_nope_query_scale(q, positions) + else: + q, k = beam_neox_rope(q, k, positions, rope_theta, 1.0) + qkv = torch.cat((q.view(num_tokens, heads * d), k.view(num_tokens, kv_heads * d), v), dim=-1) + core = torch.empty(num_tokens, heads * d, dtype=h.dtype, device=h.device) + core = thop_attention( + q=qkv, + output=core, + num_heads=heads, + num_kv_heads=kv_heads, + head_size=d, + attention_window_size=window, + **step, + **CALL_CONSTANTS, + ) + gated = beam_output_gate( + core.view(num_tokens, heads, d), gate.view(num_tokens, heads, d), gate_bias + ) + return cublas_mm(gated.view(num_tokens, heads * d), o_proj.t()) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/__init__.py new file mode 100644 index 000000000000..c29ad560f763 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam real-checkpoint loading and four-rank layout (attention TP4, MoE TP4) / sm_103 / tp4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/layout.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/layout.py new file mode 100644 index 000000000000..5376adc06d1c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/layout.py @@ -0,0 +1,369 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Configuration, schedule and four-rank ownership of the Beam checkpoint layout component. + +The layout is one fixed four-rank configuration: attention TP4 and MoE TP4 (EP1) on the +same four ranks, PP1, no attention data parallelism. Every logical checkpoint tensor is +either split along one axis into four contiguous, gap-free, non-overlapping intervals in +global checkpoint coordinates, or replicated identically on every rank:: + + embed_tokens, lm_head [151936, 5120] vocabulary rows [37984 r, 37984 (r + 1)) + q_proj, self_attn.gate_proj [6144, 5120] query-head rows [1536 r, 1536 (r + 1)) + k_proj, v_proj [1024, 5120] kv-head rows [256 r, 256 (r + 1)) + o_proj [5120, 6144] query-head columns [1536 r, 1536 (r + 1)) + gate_bias [48] query heads [12 r, 12 (r + 1)) + dense gate_proj, up_proj [20480, 5120] intermediate rows [5120 r, 5120 (r + 1)) + dense down_proj [5120, 20480] intermediate cols [5120 r, 5120 (r + 1)) + expert gate_proj, up_proj [2560, 5120] intermediate rows [640 r, 640 (r + 1)) + expert down_proj [5120, 2560] intermediate cols [640 r, 640 (r + 1)) + norms, q_norm, k_norm replicated BF16 + router weight, correction replicated FP32; expert ids stay global + +Rank r therefore owns query heads [12 r, 12 r + 12) with their two kv heads [2 r, 2 r + 2), +keeping the 6:1 grouped-query ratio, and an intermediate slice of every expert. A row-split +FP8 matrix takes the matching rows of its per-output-channel scale; a column-split one keeps +all of its output scales. Each sparse layer stores its experts' slices in three contiguous +banks, ``gate`` and ``up`` ``[256, 640, 5120]`` and ``down`` ``[256, 5120, 640]``; expert e +is slot e of each bank. +""" + +from __future__ import annotations + +import dataclasses +from collections import OrderedDict +from typing import Any, Mapping, Optional + +WORLD_SIZE = 4 + +#: The configuration this component computes, as the checkpoint's config.json states it. +SUPPORTED_CONFIG = { + "model_type": "beam", + "hidden_size": 5120, + "intermediate_size": 20480, + "moe_intermediate_size": 2560, + "num_attention_heads": 48, + "num_key_value_heads": 8, + "head_dim": 128, + "num_experts": 256, + "num_experts_per_tok": 8, + "n_group": 8, + "topk_group": 4, + "norm_topk_prob": True, + "vocab_size": 151936, + "num_hidden_layers": 52, + "rms_norm_eps": 1e-06, + "routed_scaling_factor": 2.5, + "final_logit_softcapping": 30.0, + "sliding_window": 512, + "hidden_act": "silu", + "tie_word_embeddings": False, + "attention_bias": False, +} +SLIDING_ROPE_THETA = 10000.0 + +#: Stored dtype names, as safetensors headers spell them. +FP8 = "F8_E4M3" +BF16 = "BF16" +F32 = "F32" + +ATTENTION_NORMS = ( + "input_layernorm", + "post_attention_layernorm", + "pre_feedforward_layernorm", + "post_feedforward_layernorm", +) +EMBED = "model.embed_tokens.weight" +FINAL_NORM = "model.norm.weight" +HEAD = "lm_head.weight" + + +@dataclasses.dataclass(frozen=True) +class Schedule: + """The per-layer attention and feed-forward kinds the checkpoint's configuration states.""" + + attention: tuple # "sliding_attention" | "full_attention", one per logical layer + mlp: tuple # "dense" | "sparse", one per logical layer + + @property + def num_layers(self) -> int: + return len(self.attention) + + def is_sparse(self, layer: int) -> bool: + return self.mlp[layer] == "sparse" + + def is_sliding(self, layer: int) -> bool: + return self.attention[layer] == "sliding_attention" + + +@dataclasses.dataclass(frozen=True) +class TensorSpec: + """One logical checkpoint tensor: global geometry, stored dtype and four-rank rule.""" + + name: str + shape: tuple + stored_dtype: str + family: str + layer: Optional[int] + expert: Optional[int] + split_dim: Optional[int] # None: replicated on every rank + + @property + def scale_name(self) -> Optional[str]: + """The consumed per-output-channel scale of an FP8 matrix, ``[out, 1]`` BF16.""" + return self.name + "_scale" if self.stored_dtype == FP8 else None + + def interval(self, rank: int, world_size: int = WORLD_SIZE) -> Optional[tuple]: + """``(dim, start, stop)`` of rank ``rank``'s slice, or None for a replicated tensor.""" + check_rank(rank, world_size) + if self.split_dim is None: + return None + extent = self.shape[self.split_dim] + if extent % world_size: + raise ValueError( + f"{self.name}: dim {self.split_dim} of {extent} does not split by {world_size}" + ) + size = extent // world_size + return self.split_dim, rank * size, (rank + 1) * size + + def local_shape(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + interval = self.interval(rank, world_size) + if interval is None: + return tuple(self.shape) + dim, start, stop = interval + shape = list(self.shape) + shape[dim] = stop - start + return tuple(shape) + + @property + def effective_dtype(self) -> str: + """BF16 for FP8 matrices (decoded and scaled) and BF16 tensors; F32 router tensors.""" + return BF16 if self.stored_dtype in (FP8, BF16) else F32 + + @property + def bank(self) -> Optional[str]: + """The sparse layer's bank holding this expert slice, or None for a parameter.""" + if self.expert is None: + return None + return f"layers.{self.layer}.experts.{self.family.removeprefix('expert_')}" + + +def check_rank(rank: int, world_size: int = WORLD_SIZE) -> None: + """Raise ValueError on any rank/world other than one of the four ranks of this layout.""" + if type(world_size) is not int or world_size != WORLD_SIZE: + raise ValueError( + f"this layout runs on exactly {WORLD_SIZE} ranks, got world_size={world_size!r}" + ) + if type(rank) is not int or not 0 <= rank < world_size: + raise ValueError(f"rank must be an int in [0, {world_size}), got {rank!r}") + + +def _get(config: Any, field: str) -> Any: + if isinstance(config, Mapping): + return config.get(field) + return getattr(config, field, None) + + +def check_config(config: Any) -> Schedule: + """Raise ValueError on the first field outside what this layout computes; return the schedule. + + ``config`` is the checkpoint's config.json as a mapping (or an object with those + attributes). The schedule comes from its ``layer_types``/``mlp_layer_types``: every layer + is sliding or full attention and dense or sparse, the dense layers precede the sparse + ones, and only sliding layers use RoPE (``rope_parameters``). + """ + for field, expected in SUPPORTED_CONFIG.items(): + actual = _get(config, field) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this layout computes {expected!r}") + layers = SUPPORTED_CONFIG["num_hidden_layers"] + attention = tuple(_get(config, "layer_types") or ()) + mlp = tuple(_get(config, "mlp_layer_types") or ()) + if len(attention) != layers or set(attention) - {"sliding_attention", "full_attention"}: + raise ValueError(f"layer_types must name {layers} sliding/full layers, got {attention!r}") + if len(mlp) != layers or set(mlp) - {"dense", "sparse"}: + raise ValueError(f"mlp_layer_types must name {layers} dense/sparse layers, got {mlp!r}") + if ( + "dense" in mlp + and "sparse" in mlp + and mlp.index("sparse") < len(mlp) - mlp[::-1].index("dense") + ): + raise ValueError("dense feed-forward layers must precede the sparse ones") + rope = _get(config, "rope_parameters") + sliding = (rope or {}).get("sliding_attention") if isinstance(rope, Mapping) else None + if ( + not isinstance(rope, Mapping) + or rope.get("full_attention") is not None + or not isinstance(sliding, Mapping) + or sliding.get("rope_type") != "default" + or sliding.get("rope_theta") != SLIDING_ROPE_THETA + ): + raise ValueError( + "rope_parameters must give sliding layers default RoPE with theta " + f"{SLIDING_ROPE_THETA} and full layers none, got {rope!r}" + ) + return Schedule(attention=attention, mlp=mlp) + + +def tensor_specs(config: Any, layers: Optional[tuple] = None) -> "OrderedDict[str, TensorSpec]": + """Every logical tensor of the selected layers and the outer tensors, by checkpoint name. + + ``layers`` selects logical layers (default: all of them); the embedding, final norm and + head are always included. Consumed FP8 scales are not logical tensors; see + ``TensorSpec.scale_name``. + """ + schedule = check_config(config) + hidden = _get(config, "hidden_size") + heads = _get(config, "num_attention_heads") + kv_heads = _get(config, "num_key_value_heads") + head_dim = _get(config, "head_dim") + vocab = _get(config, "vocab_size") + dense_inter = _get(config, "intermediate_size") + expert_inter = _get(config, "moe_intermediate_size") + experts = _get(config, "num_experts") + selected = tuple(range(schedule.num_layers)) if layers is None else tuple(layers) + for layer in selected: + if type(layer) is not int or not 0 <= layer < schedule.num_layers: + raise ValueError(f"layer {layer!r} is not a logical layer of the checkpoint") + if len(set(selected)) != len(selected): + raise ValueError(f"layers must be distinct, got {selected}") + + specs: "OrderedDict[str, TensorSpec]" = OrderedDict() + + def add(name, shape, dtype, family, layer=None, expert=None, split_dim=None): + specs[name] = TensorSpec(name, tuple(shape), dtype, family, layer, expert, split_dim) + + add(EMBED, (vocab, hidden), BF16, "embed_tokens", split_dim=0) + for layer in sorted(selected): + p = f"model.layers.{layer}." + for norm in ATTENTION_NORMS: + add(p + norm + ".weight", (hidden,), BF16, "block_norm", layer) + add( + p + "self_attn.q_proj.weight", + (heads * head_dim, hidden), + FP8, + "q_proj", + layer, + split_dim=0, + ) + add( + p + "self_attn.k_proj.weight", + (kv_heads * head_dim, hidden), + FP8, + "k_proj", + layer, + split_dim=0, + ) + add( + p + "self_attn.v_proj.weight", + (kv_heads * head_dim, hidden), + FP8, + "v_proj", + layer, + split_dim=0, + ) + add( + p + "self_attn.gate_proj.weight", + (heads * head_dim, hidden), + FP8, + "attn_gate", + layer, + split_dim=0, + ) + add( + p + "self_attn.o_proj.weight", + (hidden, heads * head_dim), + FP8, + "o_proj", + layer, + split_dim=1, + ) + add(p + "self_attn.q_norm.weight", (head_dim,), BF16, "qk_norm", layer) + add(p + "self_attn.k_norm.weight", (head_dim,), BF16, "qk_norm", layer) + add(p + "self_attn.gate_bias", (heads,), BF16, "gate_bias", layer, split_dim=0) + if schedule.is_sparse(layer): + add(p + "mlp.gate.weight", (experts, hidden), F32, "router_weight", layer) + add(p + "mlp.gate.e_score_correction_bias", (experts,), F32, "router_bias", layer) + for e in range(experts): + q = f"{p}mlp.experts.{e}." + add(q + "gate_proj.weight", (expert_inter, hidden), FP8, "expert_gate", layer, e, 0) + add(q + "up_proj.weight", (expert_inter, hidden), FP8, "expert_up", layer, e, 0) + add(q + "down_proj.weight", (hidden, expert_inter), FP8, "expert_down", layer, e, 1) + else: + add( + p + "mlp.gate_proj.weight", + (dense_inter, hidden), + FP8, + "dense_gate", + layer, + split_dim=0, + ) + add( + p + "mlp.up_proj.weight", (dense_inter, hidden), FP8, "dense_up", layer, split_dim=0 + ) + add( + p + "mlp.down_proj.weight", + (hidden, dense_inter), + FP8, + "dense_down", + layer, + split_dim=1, + ) + add(FINAL_NORM, (hidden,), BF16, "final_norm") + add(HEAD, (vocab, hidden), BF16, "lm_head", split_dim=0) + return specs + + +def ownership_ledger(specs: Mapping[str, TensorSpec], world_size: int = WORLD_SIZE) -> list: + """One row per logical tensor: its rule and every rank's interval, replica flag and storage. + + Rows are plain data (the same fields an independent ledger records), so two ledgers can + be compared without sharing code: ``dim``/``start``/``stop`` per split rank, ``replicated`` + otherwise, and ``bank``/``slot`` for expert slices. + """ + rows = [] + for spec in specs.values(): + ranks = [] + for rank in range(world_size): + interval = spec.interval(rank, world_size) + entry = {"rank": rank, "local_shape": list(spec.local_shape(rank, world_size))} + if interval is None: + entry["replicated"] = True + else: + dim, start, stop = interval + entry.update(replicated=False, dim=dim, start=start, stop=stop) + if spec.bank is not None: + entry.update(bank=spec.bank, slot=spec.expert) + ranks.append(entry) + rows.append( + { + "name": spec.name, + "family": spec.family, + "layer": spec.layer, + "expert": spec.expert, + "shape": list(spec.shape), + "stored_dtype": spec.stored_dtype, + "effective_dtype": spec.effective_dtype, + "scale": spec.scale_name, + "ranks": ranks, + } + ) + return rows + + +def check_partition(specs: Mapping[str, TensorSpec], world_size: int = WORLD_SIZE) -> None: + """Raise unless every split tensor's rank intervals tile its axis exactly, in rank order.""" + for spec in specs.values(): + if spec.split_dim is None: + continue + cursor = 0 + for rank in range(world_size): + dim, start, stop = spec.interval(rank, world_size) + if start != cursor or stop <= start: + raise ValueError( + f"{spec.name}: rank {rank} interval [{start}, {stop}) leaves a gap or overlap" + ) + cursor = stop + if cursor != spec.shape[spec.split_dim]: + raise ValueError( + f"{spec.name}: rank intervals cover {cursor} of {spec.shape[spec.split_dim]}" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/loader.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/loader.py new file mode 100644 index 000000000000..e03af89e1678 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/loader.py @@ -0,0 +1,338 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Metadata catalog and bounded streaming reads of a safetensors checkpoint directory. + +``SafetensorsCatalog`` reads ``model.safetensors.index.json`` and every shard's header (an +8-byte little-endian length and a JSON table); it never reads a payload. ``validate`` checks +the catalog against the expected logical tensors before anything is loaded: every expected +tensor and every FP8 matrix's scale present, nothing else present, stored dtypes and shapes +exact, each entry's byte span inside its file and matching its dtype and shape. + +``read_slice_into`` copies one rank-local slice of one tensor into its final device +storage through bounded byte-range tiles: a row-split or replicated tensor reads only its +rows; a column-split tensor reads whole rows, tile by tile, and keeps its columns on the +device. Tiles are at most ``TILE_BYTES`` and staged in at most two pinned host buffers, so +no tensor, shard or model is ever materialized on the host. An FP8 matrix's effective value +is ``bf16(float32(e4m3(q)) * float32(scale[row, 0]))``: the stored BF16 per-output-channel +scale multiplies its row, broadcast along the input columns, and the product is rounded to +BF16 once. BF16 and F32 tensors are copied byte for byte. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import struct +import threading +from collections import OrderedDict +from typing import Mapping, Optional + +import torch + +from .layout import BF16, F32, FP8, TensorSpec + +INDEX_FILE = "model.safetensors.index.json" +TILE_BYTES = 16 << 20 +_ITEMSIZE = {FP8: 1, BF16: 2, F32: 4} +_TORCH_DTYPE = {FP8: torch.float8_e4m3fn, BF16: torch.bfloat16, F32: torch.float32} +_MAX_HEADER = 256 << 20 + + +class CheckpointFormatError(ValueError): + """The checkpoint's index, headers or tensors are not what the layout expects.""" + + +class Entry: + """One stored tensor: file, dtype, shape and absolute byte span.""" + + __slots__ = ("name", "path", "dtype", "shape", "begin", "end") + + def __init__(self, name, path, dtype, shape, begin, end): + self.name, self.path, self.dtype, self.shape = name, path, dtype, tuple(shape) + self.begin, self.end = begin, end + + @property + def row_bytes(self) -> int: + cols = 1 + for d in self.shape[1:]: + cols *= d + return cols * _ITEMSIZE[self.dtype] + + +def _read_header(path: str) -> tuple: + with open(path, "rb") as f: + raw = f.read(8) + if len(raw) != 8: + raise CheckpointFormatError(f"{path}: truncated safetensors header length") + (length,) = struct.unpack(" size: + raise CheckpointFormatError( + f"{filename}: {name} spans [{begin}, {end}) outside the file" + ) + spans.append((begin, end, name)) + self.entries[name] = Entry( + name, path, dtype, shape, data_start + begin, data_start + end + ) + spans.sort() + for (b0, e0, n0), (b1, _, n1) in zip(spans, spans[1:]): + if b1 < e0: + raise CheckpointFormatError(f"{filename}: {n0} and {n1} overlap") + missing = sorted(set(weight_map) - set(self.entries)) + if missing: + raise CheckpointFormatError( + f"index entries absent from their shard headers: {missing[:8]}" + ) + + def validate(self, specs: Mapping[str, TensorSpec], *, complete: bool = True) -> None: + """Raise CheckpointFormatError unless the stored tensors are exactly what ``specs`` needs. + + With ``complete`` (``specs`` describes the whole model) any stored tensor that is + neither an expected logical tensor nor its FP8 scale is rejected; otherwise only the + expected ones are checked. + """ + expected = {} + for spec in specs.values(): + expected[spec.name] = (spec.stored_dtype, tuple(spec.shape)) + if spec.scale_name is not None: + expected[spec.scale_name] = (BF16, (spec.shape[0], 1)) + missing = sorted(set(expected) - set(self.entries)) + if missing: + raise CheckpointFormatError( + f"{len(missing)} expected tensors are missing: {missing[:8]}" + ) + if complete: + extra = sorted(set(self.entries) - set(expected)) + if extra: + raise CheckpointFormatError(f"{len(extra)} unsupported tensors: {extra[:8]}") + for name, (dtype, shape) in expected.items(): + entry = self.entries[name] + if entry.dtype != dtype or entry.shape != shape: + raise CheckpointFormatError( + f"{name} is stored as {entry.dtype}{list(entry.shape)}; expected {dtype}{list(shape)}" + ) + numel = 1 + for d in shape: + numel *= d + if entry.end - entry.begin != numel * _ITEMSIZE[dtype]: + raise CheckpointFormatError( + f"{name}: byte span does not match {dtype}{list(shape)}" + ) + + +class TileReader: + """Bounded byte-range reads: at most two pinned staging buffers of ``tile_bytes``. + + ``tiles(entry, rows)`` yields ``(row_start, row_stop, host_view)`` for consecutive row + ranges; the next tile is read on a background thread while the caller consumes the + current one, and a buffer is refilled only after the caller's device copy of it + completed. Counts the bytes read and the peak staging footprint. + """ + + def __init__(self, tile_bytes: int = TILE_BYTES, pin: bool = True): + self.tile_bytes = tile_bytes + self.buffers = [ + torch.empty(tile_bytes, dtype=torch.uint8, pin_memory=pin) for _ in range(2) + ] + self.bytes_read = 0 + self.reads = 0 + self._fds: "OrderedDict[str, int]" = OrderedDict() + self._lock = threading.Lock() + + def close(self) -> None: + for fd in self._fds.values(): + os.close(fd) + self._fds.clear() + self.buffers = [] + + def _fd(self, path: str) -> int: + fd = self._fds.get(path) + if fd is None: + if len(self._fds) >= 4: + _, old = self._fds.popitem(last=False) + os.close(old) + fd = os.open(path, os.O_RDONLY) + self._fds[path] = fd + return fd + + def _read(self, buf: torch.Tensor, fd: int, offset: int, length: int) -> None: + view = memoryview(buf.numpy())[:length] + got = 0 + while got < length: + n = os.preadv(fd, [view[got:]], offset + got) + if n <= 0: + raise OSError(f"short read at offset {offset + got}") + got += n + with self._lock: + self.bytes_read += length + self.reads += 1 + + def tiles(self, entry: Entry, row_start: int, row_stop: int): + row_bytes = entry.row_bytes + if row_bytes > self.tile_bytes: + raise CheckpointFormatError(f"{entry.name}: one row ({row_bytes} B) exceeds the tile") + rows_per_tile = max(1, self.tile_bytes // row_bytes) + fd = self._fd(entry.path) + plan = [ + (r, min(r + rows_per_tile, row_stop)) for r in range(row_start, row_stop, rows_per_tile) + ] + if not plan: + return + pending = None + try: + for i, (r0, r1) in enumerate(plan): + buf = self.buffers[i % 2] + length = (r1 - r0) * row_bytes + if pending is None: + self._read(buf, fd, entry.begin + r0 * row_bytes, length) + else: + pending.join() + if pending.error is not None: + raise pending.error + pending = None + # The other buffer held tile i - 1, whose device copy the caller completed + # before asking for tile i: read tile i + 1 into it while tile i is consumed. + if i + 1 < len(plan): + n0, n1 = plan[i + 1] + pending = _Prefetch( + self, + self.buffers[(i + 1) % 2], + fd, + entry.begin + n0 * row_bytes, + (n1 - n0) * row_bytes, + ) + pending.start() + yield r0, r1, buf[:length] + finally: + if pending is not None: + pending.join() + + +class _Prefetch(threading.Thread): + def __init__(self, reader, buf, fd, offset, length): + super().__init__(daemon=True) + self.args = (buf, fd, offset, length) + self.reader = reader + self.error = None + + def run(self): + try: + self.reader._read(*self.args) + except OSError as exc: # surfaced to the consumer by tiles() + self.error = exc + + +def _decode_tile(raw: torch.Tensor, entry: Entry, rows: int, device) -> torch.Tensor: + """Host tile bytes -> device tensor ``[rows, ...]`` in the stored dtype.""" + dev = raw.to(device, non_blocking=True) + return dev.view(_TORCH_DTYPE[entry.dtype]).view(rows, *entry.shape[1:]) + + +@torch.no_grad() +def read_slice_into( + catalog: SafetensorsCatalog, + reader: TileReader, + spec: TensorSpec, + rank: int, + dst: torch.Tensor, + *, + digest: Optional["hashlib._Hash"] = None, +) -> None: + """Copy rank ``rank``'s effective slice of ``spec`` into ``dst`` (its final device storage). + + ``dst`` must have the rank-local shape and the effective dtype. ``digest`` (optional) is + updated with the source bytes actually read, in tile order, with each tile's row range. + """ + entry = catalog.entries[spec.name] + local = spec.local_shape(rank) + want = torch.bfloat16 if spec.stored_dtype in (FP8, BF16) else torch.float32 + if tuple(dst.shape) != local or dst.dtype != want: + raise ValueError( + f"{spec.name}: destination {dst.dtype}{list(dst.shape)} != {want}{list(local)}" + ) + interval = spec.interval(rank) + rows = spec.shape[0] + row_start, row_stop, col = 0, rows, None + if interval is not None: + dim, start, stop = interval + if dim == 0: + row_start, row_stop = start, stop + else: + col = (start, stop) + scale = None + if spec.stored_dtype == FP8: + scale_entry = catalog.entries[spec.scale_name] + scale = torch.empty(row_stop - row_start, 1, dtype=torch.bfloat16, device=dst.device) + for r0, r1, raw in reader.tiles(scale_entry, row_start, row_stop): + scale[r0 - row_start : r1 - row_start].copy_( + _decode_tile(raw, scale_entry, r1 - r0, dst.device) + ) + torch.cuda.current_stream(dst.device).synchronize() + if digest is not None: + digest.update(f"{scale_entry.name}:{r0}:{r1}|".encode()) + digest.update(memoryview(raw.numpy())) + for r0, r1, raw in reader.tiles(entry, row_start, row_stop): + tile = _decode_tile(raw, entry, r1 - r0, dst.device) + if col is not None: + tile = tile[:, col[0] : col[1]] + out = dst[r0 - row_start : r1 - row_start] + if spec.stored_dtype == FP8: + s = scale[r0 - row_start : r1 - row_start].float() + out.copy_((tile.float() * s).to(torch.bfloat16)) + else: + out.copy_(tile) + torch.cuda.current_stream(dst.device).synchronize() + if digest is not None: + digest.update(f"{entry.name}:{r0}:{r1}|".encode()) + digest.update(memoryview(raw.numpy())) + del scale diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/model.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/model.py new file mode 100644 index 000000000000..b76bda2ca58c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/model.py @@ -0,0 +1,717 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam checkpoint layout / sm_103 / tp4: rank-local real weights and representative blocks. + +A component, not a routed target. There is no ``routing.py`` and no ``ForCausalLM``: a +caller on each of four ranks constructs this class with the checkpoint's configuration and +its rank's ``Mapping`` (attention TP4, MoE TP4/EP1, PP1), loads a caller-supplied checkpoint +directory into the rank's slices, and runs representative computations. It does not run a +full-depth forward, prefill/decode loop or sampler. + +Storage is rank-local and final: each logical tensor's slice (``layout``) lives in its own +parameter, except that each sparse layer keeps its 256 experts' slices in three contiguous +banks. Loading streams bounded tiles straight into that storage (``loader``) and is a +transaction: the component is usable only after every rank loaded every selected tensor. + +Per call, rank r computes its 12 query heads / 2 kv heads, its dense or expert intermediate +slice and its vocabulary rows; the enclosing method applies exactly one collective at each +boundary, on the TRT-LLM NCCL all-reduce or all-gather:: + + embed masked local lookup, one all-reduce (sum) + attention local q/k/v/gate, q/k norms, RoPE or NoPE scale, TRTLLM attention on the + rank's kv heads, output gate, local o_proj partial, one all-reduce + dense FFN local gate/up, fp32 SwiGLU, local down partial, one all-reduce + MoE replicated fp32 router (global expert ids), all-rank agreement on the + router input, ids and weights, rank-local expert slices combined in + ascending global expert order (bf16 routes, fp32 sum, x2.5, bf16), one + all-reduce + head final norm, local vocabulary logits, one all-gather along the + vocabulary in rank order, fp32 softcap once + +Every method accepts ``capture``, a dict that receives the named intermediates of the call +(clones); with ``capture=None`` nothing is recorded. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from collections import OrderedDict +from typing import Any, Mapping, Optional, Sequence + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_neox_rope import ( + beam_neox_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_nope_query_scale import ( # noqa: E501 + beam_nope_query_scale, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_output_gate import ( + beam_output_gate, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + thop_attention, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.residual.beam_depth_residual import ( + beam_depth_residual, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.distributed import AllReduce, allgather +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.mapping import Mapping as ParallelMapping + +from .. import _attention +from ..dense_block_bf16__sm_103__tp1.block import depth_scales +from ..moe_dispatch_combine_bf16__sm_103__tp1 import dispatch_combine +from ..moe_expert_bf16__sm_103__tp1.ops.activation.swiglu_fp32 import swiglu_fp32 +from ..moe_router_bf16__sm_103__tp1.ops.moe.beam_grouped_topk_router import beam_grouped_topk_router +from ..small_model_prefill_bf16__sm_103__tp1.ops.activation.beam_logit_softcap import ( + beam_logit_softcap, +) +from .layout import ( + EMBED, + FINAL_NORM, + HEAD, + WORLD_SIZE, + TensorSpec, + check_config, + check_partition, + check_rank, + ownership_ledger, + tensor_specs, +) +from .loader import CheckpointFormatError, SafetensorsCatalog, TileReader, read_slice_into +from .ops.norm.beam_rmsnorm_fp32 import beam_rmsnorm_fp32 + +_SM = (10, 3) +ROUTER_EPS = 1e-20 +#: Default tokens per call; the representative cases use at most 129. A caller that runs longer +#: calls (a full-model prefill) passes its own ``max_tokens`` at construction. +MAX_TOKENS = 1024 + + +class LoadError(RuntimeError): + """A load failed on at least one rank; the component is not usable on any rank.""" + + +def check_max_tokens(max_tokens: int) -> int: + """Return ``max_tokens`` if it is a positive int (the most tokens one call accepts).""" + if type(max_tokens) is not int or max_tokens < 1: + raise ValueError(f"max_tokens must be a positive int, got {max_tokens!r}") + return max_tokens + + +def _check_mapping(mapping: ParallelMapping) -> None: + check_rank(mapping.rank, mapping.world_size) + expected = dict(tp_size=4, pp_size=1, moe_tp_size=4, moe_ep_size=1, cp_size=1) + for field, value in expected.items(): + actual = getattr(mapping, field) + if actual != value: + raise ValueError(f"this layout needs Mapping {field}={value}, got {actual}") + if mapping.enable_attention_dp: + raise ValueError("this layout does not use attention data parallelism") + if list(mapping.tp_group) != list(range(WORLD_SIZE)): + raise ValueError(f"the TP group must be ranks 0..3, got {mapping.tp_group}") + + +def _tensor_state(tensor: torch.Tensor) -> tuple: + return ( + id(tensor), + tensor._version, + tensor.data_ptr(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device, + ) + + +class BeamCheckpointLayoutBf16Sm103Tp4(nn.Module): + """Rank-local real-checkpoint weights of the Beam model on one of four TP ranks. + + ``config`` is the checkpoint's config.json (a mapping); ``mapping`` this rank's + ``tensorrt_llm.mapping.Mapping`` with world size 4, TP4, PP1, MoE TP4/EP1; ``layers`` + the logical layers to hold (default all 52; the embedding, final norm and head are + always held); ``max_tokens`` the most tokens one call accepts (default ``MAX_TOKENS``). + Construction allocates uninitialized rank-local storage on the current CUDA device and + validates the configuration and layout; it reads no checkpoint data. + """ + + def __init__( + self, + config: Mapping[str, Any], + mapping: ParallelMapping, + layers: Optional[Sequence[int]] = None, + device: Optional[torch.device] = None, + *, + max_tokens: int = MAX_TOKENS, + ): + super().__init__() + self.max_tokens = check_max_tokens(max_tokens) + self.schedule = check_config(config) + _check_mapping(mapping) + device = ( + torch.device("cuda", torch.cuda.current_device()) + if device is None + else torch.device(device) + ) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this component is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.config = dict(config) + self.mapping = mapping + self.rank = mapping.rank + self.device = device + self.layers = tuple(sorted(range(self.schedule.num_layers) if layers is None else layers)) + self.specs = tensor_specs(config, self.layers) + check_partition(self.specs) + self.hidden = config["hidden_size"] + self.head_dim = config["head_dim"] + self.local_heads = config["num_attention_heads"] // WORLD_SIZE + self.local_kv_heads = config["num_key_value_heads"] // WORLD_SIZE + self.eps = config["rms_norm_eps"] + self.window = config["sliding_window"] + self.rope_theta = config["rope_parameters"]["sliding_attention"]["rope_theta"] + self.softcap = config["final_logit_softcapping"] + self.routed_scaling_factor = config["routed_scaling_factor"] + self.vocab_interval = self.specs[EMBED].interval(self.rank)[1:] + #: Storage is always resolved from the currently registered parameters, never from + #: references taken here, so a replaced parameter or bank is what every call uses. + self.tensor_names: list = [] + self.bank_names: list = [] + for spec in self.specs.values(): + if spec.bank is not None: + if spec.bank not in self.bank_names: + local = spec.local_shape(self.rank) + bank = torch.empty( + (config["num_experts"], *local), dtype=torch.bfloat16, device=device + ) + self.register_parameter( + _attr(spec.bank), nn.Parameter(bank, requires_grad=False) + ) + self.bank_names.append(spec.bank) + continue + dtype = torch.bfloat16 if spec.effective_dtype == "BF16" else torch.float32 + param = nn.Parameter( + torch.empty(spec.local_shape(self.rank), dtype=dtype, device=device), + requires_grad=False, + ) + self.register_parameter(_attr(spec.name), param) + self.tensor_names.append(spec.name) + self._allreduce = AllReduce(mapping=mapping, strategy=AllReduceStrategy.NCCL) + self.usable = False + self.load_record: Optional[dict] = None + + # ------------------------------------------------------------------ storage + + def tensor(self, name: str) -> torch.Tensor: + """The rank-local storage of logical tensor ``name`` (an expert's view of its bank slot).""" + spec = self.specs[name] + if spec.bank is not None: + return self.bank(spec.bank)[spec.expert] + return getattr(self, _attr(name)) + + def bank(self, bank: str) -> torch.Tensor: + """The current ``[256, ...]`` storage of one sparse layer's expert bank.""" + return getattr(self, _attr(bank)) + + def ledger(self) -> list: + """This component's ownership rows (see ``layout.ownership_ledger``).""" + return ownership_ledger(self.specs) + + def weight_state(self) -> tuple: + """Identity of every current parameter and bank: object, version, storage, geometry. + + Equal states mean no parameter was replaced or modified in place since the state + was taken; a holder of derived state (a cache) binds it and checks it before use. + """ + return tuple((name, _tensor_state(p)) for name, p in self.named_parameters()) + + def check_weight_state(self, state: tuple) -> None: + """Raise RuntimeError unless the weights are exactly those ``state`` was taken from.""" + if not self.usable: + raise RuntimeError("the component is not usable: no complete load") + if self.weight_state() != state: + raise RuntimeError( + "the weights changed after the state was bound (in place or by replacement); " + "derived state such as a KV cache no longer matches them" + ) + + # ------------------------------------------------------------------ loading + + def _agree(self, ok: bool) -> bool: + """All-rank AND of ``ok`` through one all-reduce; every rank must call it.""" + flag = torch.tensor([0.0 if ok else 1.0], dtype=torch.float32, device=self.device) + return float(self._allreduce(flag).item()) == 0.0 + + @torch.no_grad() + def load(self, checkpoint_dir: str, *, record_digests: bool = True) -> dict: + """Stream the selected tensors of a caller-supplied checkpoint into rank-local storage. + + Collective: every rank calls it. Metadata is validated against the whole model + (every expected tensor and scale, nothing unsupported, exact dtypes and shapes) and + the directory's config.json against the constructed configuration before any + storage changes. The component is unusable from the first write until every rank + reports a complete load; a failure on any rank raises ``LoadError`` on all ranks + and leaves the component unusable. Returns the load record (coverage, source + digests, bytes read). + """ + error: Optional[BaseException] = None + catalog = None + try: + with open(os.path.join(checkpoint_dir, "config.json")) as f: + stored = json.load(f) + if stored != self.config: + differing = sorted( + k for k in set(stored) | set(self.config) if stored.get(k) != self.config.get(k) + ) + raise CheckpointFormatError( + f"config.json differs from the constructed configuration: {differing[:8]}" + ) + catalog = SafetensorsCatalog(checkpoint_dir) + catalog.validate(tensor_specs(self.config), complete=True) + except Exception as exc: # Keep rank-local failures inside the collective protocol. + error = exc + if not self._agree(error is None): + raise LoadError( + f"rank {self.rank}: checkpoint rejected before loading: {error!r}" + ) from error + + self.usable = False + self.load_record = None + coverage = OrderedDict() + reader = None + bytes_read = reads = 0 + try: + reader = TileReader() + order = sorted( + self.specs.values(), + key=lambda s: (catalog.entries[s.name].path, catalog.entries[s.name].begin), + ) + for spec in order: + digest = hashlib.sha256() if record_digests else None + dst = self.tensor(spec.name) + read_slice_into(catalog, reader, spec, self.rank, dst, digest=digest) + coverage[spec.name] = { + "interval": spec.interval(self.rank), + "local_shape": list(dst.shape), + "dtype": str(dst.dtype), + "source_sha256": digest.hexdigest() if digest is not None else None, + "bank": spec.bank, + "slot": spec.expert, + } + except Exception as exc: # Every rank must reach the completion agreement. + error = exc + finally: + if reader is not None: + bytes_read, reads = reader.bytes_read, reader.reads + try: + reader.close() + except Exception as exc: # Cleanup failure must also reach the agreement. + if error is None: + error = exc + complete = error is None and list(coverage) == [s.name for s in order] + if not self._agree(complete): + raise LoadError( + f"rank {self.rank}: load did not complete on every rank: {error!r}" + ) from error + self.load_record = { + "rank": self.rank, + "checkpoint_dir": os.fspath(checkpoint_dir), + "layers": list(self.layers), + "tensors": len(coverage), + "bytes_read": bytes_read, + "reads": reads, + "coverage": coverage, + } + self.usable = True + return self.load_record + + def effective_sha256(self, name: str) -> str: + """SHA-256 of the rank-local effective bytes of ``name`` (C order, its dtype).""" + t = self.tensor(name).contiguous().cpu() + return hashlib.sha256(memoryview(t.view(torch.uint8).numpy())).hexdigest() + + # ------------------------------------------------------------------ helpers + + def _check_usable(self, layer: Optional[int] = None) -> None: + if not self.usable: + raise RuntimeError("the component is not usable: no complete load on every rank") + if layer is not None and layer not in self.layers: + raise ValueError(f"logical layer {layer} is not held by this component ({self.layers})") + + def _layer(self, layer: int, name: str) -> torch.Tensor: + return self.tensor(f"model.layers.{layer}.{name}") + + def _bank(self, layer: int, which: str) -> torch.Tensor: + return self.bank(f"layers.{layer}.experts.{which}") + + def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor: + """The one sum over the four ranks at a partial-result boundary (NCCL).""" + return self._allreduce(tensor) + + def _digests_agree(self, tensors: Sequence[torch.Tensor]) -> list: + """Raise unless every rank holds byte-identical ``tensors``; return this rank's digests.""" + digests = [] + for t in tensors: + h = hashlib.sha256(f"{t.dtype}{list(t.shape)}".encode()) + h.update(memoryview(t.contiguous().cpu().view(torch.uint8).numpy().reshape(-1))) + digests.append(h.digest()) + local = torch.frombuffer(bytearray(b"".join(digests)), dtype=torch.uint8).to(self.device) + gathered = allgather(local, self.mapping, dim=0).view(WORLD_SIZE, -1).cpu() + if not all(torch.equal(gathered[0], gathered[r]) for r in range(1, WORLD_SIZE)): + raise RuntimeError( + "replicated MoE routing differs between ranks (input, ids or weights)" + ) + return [d.hex() for d in digests] + + # ------------------------------------------------------------------ computations + + @torch.no_grad() + def embed(self, token_ids: torch.Tensor, capture: Optional[dict] = None) -> torch.Tensor: + """Global token ids ``[T]`` int64 -> replicated embeddings ``[T, hidden]`` bf16.""" + self._check_usable() + _check_ids(token_ids, self.config["vocab_size"], self.device, self.max_tokens) + lo, hi = self.vocab_interval + inside = (token_ids >= lo) & (token_ids < hi) + local = nn.functional.embedding( + torch.where(inside, token_ids - lo, torch.zeros_like(token_ids)), self.tensor(EMBED) + ) + part = torch.where(inside.unsqueeze(-1), local, torch.zeros_like(local)) + out = self.all_reduce(part) + _rec(capture, {"tp.embed_part": part, "embed": out}) + return out + + @torch.no_grad() + def attention( + self, + layer: int, + h: torch.Tensor, + positions: torch.Tensor, + md: TrtllmAttentionMetadata, + capture: Optional[dict] = None, + *, + generation_rows: int = 1, + extra_kv_tokens: int = 0, + ) -> torch.Tensor: + """Normalized input ``[T, hidden]`` -> attention output after the one o_proj all-reduce. + + ``md`` is one prepared step of one or more requests (``_attention.check_batch_metadata`` + with ``generation_rows`` rows per generation request and ``extra_kv_tokens`` reserved + KV slots) whose new tokens are the T rows, contexts first; ``positions`` are their + absolute positions. + """ + self._check_usable(layer) + _check_hidden(h, self.hidden, self.device, self.max_tokens) + T = h.shape[0] + _attention.check_batch_metadata( + md, + T, + layer_idx=layer, + num_kv_heads=self.local_kv_heads, + head_dim=self.head_dim, + generation_rows=generation_rows, + extra_kv_tokens=extra_kv_tokens, + ) + step = _attention.step_args(md, layer) + heads, kv_heads, d = self.local_heads, self.local_kv_heads, self.head_dim + q = cublas_mm(h, self._layer(layer, "self_attn.q_proj.weight").t()) + k = cublas_mm(h, self._layer(layer, "self_attn.k_proj.weight").t()) + v = cublas_mm(h, self._layer(layer, "self_attn.v_proj.weight").t()) + g = cublas_mm(h, self._layer(layer, "self_attn.gate_proj.weight").t()) + qn = flashinfer_rmsnorm( + q.view(T * heads, d), self._layer(layer, "self_attn.q_norm.weight"), self.eps + ) + kn = flashinfer_rmsnorm( + k.view(T * kv_heads, d), self._layer(layer, "self_attn.k_norm.weight"), self.eps + ) + qn, kn = qn.view(T, heads, d), kn.view(T, kv_heads, d) + sliding = self.schedule.is_sliding(layer) + if sliding: + qa, ka = beam_neox_rope(qn, kn, positions, self.rope_theta, 1.0) + else: + qa, ka = beam_nope_query_scale(qn, positions), kn + qkv = torch.cat((qa.reshape(T, heads * d), ka.reshape(T, kv_heads * d), v), dim=-1) + core = torch.empty(T, heads * d, dtype=h.dtype, device=h.device) + core = thop_attention( + q=qkv, + output=core, + num_heads=heads, + num_kv_heads=kv_heads, + head_size=d, + attention_window_size=self.window if sliding else step["max_seq_len"], + **step, + **_attention.call_constants(generation_rows), + ) + gated = beam_output_gate( + core.view(T, heads, d), g.view(T, heads, d), self._layer(layer, "self_attn.gate_bias") + ) + part = cublas_mm( + gated.reshape(T, heads * d), self._layer(layer, "self_attn.o_proj.weight").t() + ) + out = self.all_reduce(part) + _rec( + capture, + { + "q_proj": q, + "k_proj": k, + "v_proj": v, + "gate_proj": g, + "q_norm": qn, + "k_norm": kn, + ("q_rope" if sliding else "q_scaled"): qa, + **({"k_rope": ka} if sliding else {}), + "q_attn": qa, + "k_attn": ka, + "attn_core": core.view(T, heads, d), + "gated": gated, + "tp.o_part": part, + "attn_out": out, + }, + ) + return out + + @torch.no_grad() + def o_proj_reduce( + self, layer: int, gated: torch.Tensor, capture: Optional[dict] = None + ) -> torch.Tensor: + """Local heads' gated attention ``[T, 12, 128]`` -> o_proj partial, reduced once.""" + self._check_usable(layer) + T = gated.shape[0] + part = cublas_mm( + gated.reshape(T, self.local_heads * self.head_dim), + self._layer(layer, "self_attn.o_proj.weight").t(), + ) + out = self.all_reduce(part) + _rec(capture, {"tp.o_part": part, "attn_out": out}) + return out + + @torch.no_grad() + def dense_ffn( + self, layer: int, h: torch.Tensor, capture: Optional[dict] = None + ) -> torch.Tensor: + """Normalized input ``[T, hidden]`` -> dense SwiGLU output after the one down all-reduce.""" + self._check_usable(layer) + _check_hidden(h, self.hidden, self.device, self.max_tokens) + if self.schedule.is_sparse(layer): + raise ValueError(f"layer {layer} has a sparse feed-forward") + gate = cublas_mm(h, self._layer(layer, "mlp.gate_proj.weight").t()) + up = cublas_mm(h, self._layer(layer, "mlp.up_proj.weight").t()) + act = swiglu_fp32(gate, up) + part = cublas_mm(act["act"], self._layer(layer, "mlp.down_proj.weight").t()) + out = self.all_reduce(part) + _rec( + capture, + { + "mlp.gate": gate, + "mlp.up": up, + "mlp.silu": act["silu"], + "mlp.product": act["product"], + "mlp.act": act["act"], + "tp.mlp_part": part, + "mlp_out": out, + }, + ) + return out + + @torch.no_grad() + def route(self, layer: int, h: torch.Tensor) -> dict: + """The replicated FP32 router on ``h`` ``[T, hidden]``: every routing intermediate.""" + self._check_usable(layer) + _check_hidden(h, self.hidden, self.device, self.max_tokens) + if not self.schedule.is_sparse(layer): + raise ValueError(f"layer {layer} has a dense feed-forward") + return beam_grouped_topk_router( + h, + self._layer(layer, "mlp.gate.weight"), + self._layer(layer, "mlp.gate.e_score_correction_bias"), + n_group=self.config["n_group"], + topk_group=self.config["topk_group"], + top_k=self.config["num_experts_per_tok"], + eps=ROUTER_EPS, + ) + + @torch.no_grad() + def moe(self, layer: int, h: torch.Tensor, capture: Optional[dict] = None) -> torch.Tensor: + """Normalized input ``[T, hidden]`` -> MoE output after the one all-reduce. + + Before any expert runs, every rank's router input bytes, ordered expert ids and + FP32 weights must be identical (one all-gather of their digests). + """ + routing = self.route(layer, h) + ids, weights = routing["ids"], routing["weights"] + digests = self._digests_agree((h, ids, weights)) + gate_bank, up_bank, down_bank = (self._bank(layer, w) for w in ("gate", "up", "down")) + calls = [] + + def local_expert(expert_id: int, rows: torch.Tensor) -> torch.Tensor: + g = cublas_mm(rows, gate_bank[expert_id].t()) + u = cublas_mm(rows, up_bank[expert_id].t()) + act = swiglu_fp32(g, u) + y = cublas_mm(act["act"], down_bank[expert_id].t()) + if capture is not None: + calls.append((expert_id, rows, g, u, act, y)) + return y + + local = dispatch_combine.beam_dispatch_combine(h, ids, weights, local_expert) + out = self.all_reduce(local) + if capture is not None: + _rec( + capture, + { + "moe.logits": routing["logits"], + "moe.scores": routing["scores"], + "moe.selection": routing["selection_scores"], + "moe.group_scores": routing["group_scores"], + "moe.groups": routing["groups"], + "moe.masked": routing["masked_selection_scores"], + "moe.ids": ids, + "moe.weights": weights, + "moe.call_expert": torch.tensor( + [e for e, rows, *_ in calls for _ in range(rows.shape[0])], + dtype=torch.int64, + ), + "moe.expert_in": torch.cat([c[1] for c in calls]), + "moe.gate": torch.cat([c[2] for c in calls]), + "moe.up": torch.cat([c[3] for c in calls]), + "moe.silu": torch.cat([c[4]["silu"] for c in calls]), + "moe.product": torch.cat([c[4]["product"] for c in calls]), + "moe.act": torch.cat([c[4]["act"] for c in calls]), + "moe.tp.y": torch.cat([c[5] for c in calls]), + "moe.tp.local_out": local, + "moe.output": out, + "mlp_out": out, + }, + ) + capture["moe.agreement_sha256"] = digests + return out + + @torch.no_grad() + def block( + self, + layer: int, + x: torch.Tensor, + positions: torch.Tensor, + md: TrtllmAttentionMetadata, + capture: Optional[dict] = None, + *, + generation_rows: int = 1, + extra_kv_tokens: int = 0, + ) -> torch.Tensor: + """One decoder block of logical layer ``layer`` on ``x`` ``[T, hidden]``, one engine step. + + The input norm runs the source model's torch op order (``beam_rmsnorm_fp32``), so the + q/k/v/gate projections read the source's normalized bytes for the same block input. + The other norms use the flashinfer kernel. ``generation_rows`` and ``extra_kv_tokens`` + are the step's speculative shape (see ``attention``). + """ + self._check_usable(layer) + p = f"model.layers.{layer}." + h = beam_rmsnorm_fp32(x, self.tensor(p + "input_layernorm.weight"), self.eps) + a = self.attention( + layer, + h, + positions, + md, + capture, + generation_rows=generation_rows, + extra_kv_tokens=extra_kv_tokens, + ) + pa = flashinfer_rmsnorm(a, self.tensor(p + "post_attention_layernorm.weight"), self.eps) + x1 = beam_depth_residual(pa, x, *depth_scales(layer, is_attention=True)) + pf = flashinfer_rmsnorm(x1, self.tensor(p + "pre_feedforward_layernorm.weight"), self.eps) + m = ( + self.moe(layer, pf, capture) + if self.schedule.is_sparse(layer) + else self.dense_ffn(layer, pf, capture) + ) + pm = flashinfer_rmsnorm(m, self.tensor(p + "post_feedforward_layernorm.weight"), self.eps) + out = beam_depth_residual(pm, x1, *depth_scales(layer, is_attention=False)) + _rec( + capture, + { + "input": x, + "input_norm": h, + "post_attn_norm": pa, + "attn_residual": x1, + "pre_ff_norm": pf, + "post_ff_norm": pm, + "block_output": out, + }, + ) + return out + + @torch.no_grad() + def final_norm(self, h: torch.Tensor, capture: Optional[dict] = None) -> torch.Tensor: + self._check_usable() + _check_hidden(h, self.hidden, self.device, self.max_tokens) + out = flashinfer_rmsnorm(h, self.tensor(FINAL_NORM), self.eps) + _rec(capture, {"final_norm": out}) + return out + + @torch.no_grad() + def head(self, n: torch.Tensor, capture: Optional[dict] = None) -> torch.Tensor: + """Final-norm output ``[T, hidden]`` -> softcapped fp32 logits over the full vocabulary. + + The rank's vocabulary rows give local logits; one all-gather along the vocabulary + axis restores global order (rank 0's rows first), and the softcap is applied once. + """ + self._check_usable() + _check_hidden(n, self.hidden, self.device, self.max_tokens) + local = cublas_mm(n, self.tensor(HEAD).t()) + gathered = allgather(local, self.mapping, dim=-1) + logits = beam_logit_softcap(gathered, self.softcap) + _rec(capture, {"tp.head_part": local, "head": gathered, "logits": logits}) + return logits + + +def _attr(name: str) -> str: + return name.replace(".", "__") + + +def _rec(capture: Optional[dict], items: dict) -> None: + if capture is None: + return + for name, tensor in items.items(): + capture[name] = tensor.detach().clone() + + +def _check_hidden( + h: torch.Tensor, hidden: int, device: torch.device, max_tokens: int = MAX_TOKENS +) -> None: + if not isinstance(h, torch.Tensor) or h.dtype != torch.bfloat16: + raise ValueError("hidden states must be a torch.bfloat16 tensor") + if h.dim() != 2 or h.shape[1] != hidden or not 1 <= h.shape[0] <= max_tokens: + raise ValueError( + f"hidden states must be [T, {hidden}] with 1 <= T <= {max_tokens}, got {list(h.shape)}" + ) + if h.device != device or not h.is_contiguous(): + raise ValueError("hidden states must be contiguous on the component's device") + + +def _check_ids( + ids: torch.Tensor, vocab: int, device: torch.device, max_tokens: int = MAX_TOKENS +) -> None: + if not isinstance(ids, torch.Tensor) or ids.dtype != torch.int64 or ids.dim() != 1: + raise ValueError("token ids must be a 1-D torch.int64 tensor") + if ids.device != device or not 1 <= ids.shape[0] <= max_tokens: + raise ValueError(f"token ids must hold 1..{max_tokens} ids on the component's device") + if bool(((ids < 0) | (ids >= vocab)).any()): + raise ValueError(f"token ids must be in [0, {vocab})") + + +__all__ = [ + "BeamCheckpointLayoutBf16Sm103Tp4", + "LoadError", + "MAX_TOKENS", + "check_max_tokens", + "TensorSpec", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/__init__.py new file mode 100644 index 000000000000..23e614d4422b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this component. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/norm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/norm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/norm/beam_rmsnorm_fp32.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/norm/beam_rmsnorm_fp32.py new file mode 100644 index 000000000000..6c1a635fc0db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/norm/beam_rmsnorm_fp32.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RMS normalization over the last dim in the source model's fp32 op order, rounded once.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The Beam decoder block's input norm: one 5120-wide row per token, bf16, eps 1e-6. +_BLOCK_INPUT = dict(width=5120, dtype=torch.bfloat16, eps=1e-6) + + +class BeamRmsnormFp32(OpWrapper): + """`x / sqrt(mean(x^2, -1) + eps) * weight`, one eager torch op per source step. + + `x` is `[T, D]` bf16 and `weight` `[D]` bf16. The steps, as the source model computes + them:: + + x32 = x.float() [T, D] fp32 + rms = x32.pow(2).mean(-1, keepdim=True) [T, 1] fp32 + .add(eps).rsqrt() + normed = x32 * rms [T, D] fp32 + out = (normed * weight).to(x.dtype) [T, D] bf16, weight promoted + + Returns `out`, a new tensor; `x` is not mutated. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block input norm, 1 token", spec=dict(_BLOCK_INPUT, tokens=1)), + Cell(why="block input norm, 17 tokens", spec=dict(_BLOCK_INPUT, tokens=17)), + Cell( + why="block input norm, 129 tokens: the longest context", + spec=dict(_BLOCK_INPUT, tokens=129), + ), + ) + + note = """ + Not the catalog `flashinfer_rmsnorm` entry. That kernel also reduces in fp32 and rounds + once, but it sums the squares in its own order and computes its own reciprocal square + root. On identical bf16 inputs its output differs from the source model's by one bf16 + unit wherever the fp32 product lands near a rounding boundary. The difference is a few + elements per 129 tokens of real hidden states. Every projection that reads the normalized + row then carries it. + + This entry issues the source's torch ops in the source's order. Its result equals the + source's when both run the same torch reduction kernel for the same row width. The + entry does not guard against a different torch build that changes that kernel's + summation order. + """ + + def __call__(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + x32 = x.float() + rms = x32.pow(2).mean(dim=-1, keepdim=True).add(eps).rsqrt() + normed = x32 * rms + return (normed * weight).to(x.dtype) + + def reference(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Statistics and products in fp64, rounded once to `x.dtype`.""" + x64 = x.double() + rms = torch.rsqrt(x64.pow(2).mean(dim=-1, keepdim=True) + eps) + return (x64 * rms * weight.double()).to(x.dtype) + + def is_valid(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> None: + # A weight of another dtype or width broadcasts or promotes without raising, and an + # fp32 input removes the one rounding the projections' input carries. + assert x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( + f"x and weight must be torch.bfloat16, got {x.dtype} and {weight.dtype}" + ) + assert weight.dim() == 1 and weight.shape[0] == x.shape[-1], ( + f"weight must be [{x.shape[-1]}], got {tuple(weight.shape)}" + ) + + +beam_rmsnorm_fp32 = BeamRmsnormFp32() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/__init__.py new file mode 100644 index 000000000000..cf69388c4f81 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam NVFP4-expert checkpoint loading and four-rank layout (attention TP4, MoE TP4) / sm_103 / tp4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/expert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/expert.py new file mode 100644 index 000000000000..c722fd896964 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/expert.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 W4A4 SwiGLU expert on one rank's intermediate slice: the MoE expert provider. + +``nvfp4_swiglu_expert`` is the provider ``beam_dispatch_combine`` calls once per selected expert, +in ascending expert order, with that expert's routed rows ``x`` ``[n, 5120]`` bf16. With the +rank's packed codes ``W``, 128x4-swizzled block scales ``S`` and FP32 scalars of expert ``e``:: + + fc1, sf1 = fp4_quantize(x, g_gate[e]) [n, 2560] u8, swizzled scales + gate = nvfp4_gemm(fc1, W_gate[e], sf1, S_gate[e], alpha_gate[e]) [n, 640] bf16 + up = nvfp4_gemm(fc1, W_up[e], sf1, S_up[e], alpha_up[e]) [n, 640] bf16 + act = bf16(silu(gate.float()) * up.float()) swiglu_fp32 + fc2, sf2 = fp4_quantize(act, g_down[e]) [n, 320] u8, swizzled scales + y = nvfp4_gemm(fc2, W_down[e], sf2, S_down[e], alpha_down[e]) [n, 5120] bf16 + +``g = 1 / input_scale`` and ``alpha = input_scale * weight_scale_2`` per projection. Gate and up +share one activation quantization only when their stored input scales are bitwise equal; +otherwise up quantizes ``x`` with its own ``g``. Every GEMM is the CUTLASS NVFP4 GEMM +(``allowed_backends="cutlass"``) with BF16 output; there is no BF16 expert path. The output is +the expert's unweighted rank partial: mixing weights, the routed scaling factor and the +tensor-parallel sum are applied by the caller. + +Per expert this is two quantizer calls (three when the gate/up input scales differ), three +GEMMs and one elementwise SwiGLU stage. Expert enumeration happens on the host in +``beam_dispatch_combine``, so the path synchronizes with the device and is not CUDA-graph +capturable; a batched pre-routed backend can replace this provider without changing routing, +combination or the reduction. +""" + +from __future__ import annotations + +from typing import Mapping, Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.nvfp4_gemm import nvfp4_gemm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) + +from ..moe_expert_bf16__sm_103__tp1.ops.activation.swiglu_fp32 import swiglu_fp32 +from .layout import GROUP + +#: The only NVFP4 GEMM backend this provider runs. +GEMM_BACKENDS = "cutlass" + +#: The kernel contract, as recorded in load and run records. +CONTRACT = { + "activation_quantizer": "torch.ops.trtllm.fp4_quantize(x, g, sf_vec_size=16, " + "sf_use_ue8m0=False, is_sf_swizzled_layout=True)", + "gemm": "torch.ops.trtllm.nvfp4_gemm(..., output_dtype=torch.bfloat16, " + f"allowed_backends={GEMM_BACKENDS!r})", + "activation": "swiglu_fp32: bf16(silu(gate.float()) * up.float())", + "g": "float32(1 / input_scale)", + "alpha": "float32(input_scale * weight_scale_2)", + "weight_scales": "E4M3 bytes, rank slice first, then torch.ops.trtllm.block_scale_interleave", +} + + +def _gemm(a: torch.Tensor, a_sf: torch.Tensor, proj: Mapping[str, torch.Tensor], e: int): + return nvfp4_gemm( + a, + proj["weight"][e], + a_sf, + proj["weight_scale"][e], + proj["alpha"][e : e + 1], + torch.bfloat16, + 0, + GEMM_BACKENDS, + ) + + +def _quantize(x: torch.Tensor, proj: Mapping[str, torch.Tensor], e: int): + return fp4_quantize(x, proj["g"][e : e + 1], GROUP, False, True) + + +@torch.no_grad() +def nvfp4_swiglu_expert( + rows: torch.Tensor, + bank: Mapping[str, Mapping[str, torch.Tensor]], + expert: int, + shared_fc1: bool, + record: Optional[list] = None, +) -> torch.Tensor: + """Expert ``expert``'s rank partial for ``rows`` ``[n, 5120]`` bf16 -> ``[n, 5120]`` bf16. + + ``bank`` maps ``"gate"``, ``"up"``, ``"down"`` to the layer's banks (``weight``, + ``weight_scale``, ``g``, ``alpha``, each indexed by global expert id); ``shared_fc1`` says + the gate and up input scales of this expert are bitwise equal. With ``record`` (a list) the + call's operands and intermediates are appended as one dict. + """ + gate_p, up_p, down_p = bank["gate"], bank["up"], bank["down"] + fc1, sf1 = _quantize(rows, gate_p, expert) + fc1_up, sf1_up = (fc1, sf1) if shared_fc1 else _quantize(rows, up_p, expert) + gate = _gemm(fc1, sf1, gate_p, expert) + up = _gemm(fc1_up, sf1_up, up_p, expert) + act = swiglu_fp32(gate, up) + fc2, sf2 = _quantize(act["act"], down_p, expert) + y = _gemm(fc2, sf2, down_p, expert) + if record is not None: + record.append( + { + "expert": expert, + "rows": rows, + "fc1": fc1, + "sf1": sf1, + "fc1_up": None if shared_fc1 else fc1_up, + "sf1_up": None if shared_fc1 else sf1_up, + "gate": gate, + "up": up, + "silu": act["silu"], + "product": act["product"], + "act": act["act"], + "fc2": fc2, + "sf2": sf2, + "y": y, + "scalars": torch.stack( + [ + gate_p["g"][expert], + up_p["g"][expert], + gate_p["alpha"][expert], + up_p["alpha"][expert], + down_p["g"][expert], + down_p["alpha"][expert], + ] + ), + } + ) + return y diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/layout.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/layout.py new file mode 100644 index 000000000000..7ecdca806246 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/layout.py @@ -0,0 +1,335 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Recipe, stored geometry and four-rank ownership of the Beam NVFP4-expert checkpoint layout. + +The checkpoint is a ModelOpt NVFP4 export (``hf_quant_config.json``: ``quant_algo`` NVFP4, +group size 16, static activation scales, no KV-cache quantization). Only the routed experts' +gate/up/down projections are quantized; every other tensor keeps its stored precision: the +attention, dense feed-forward, norm, embedding and head tensors are BF16, the router's weight and +correction bias FP32. The four-rank configuration and every logical tensor's interval are those +of the BF16 layout (``checkpoint_layout_bf16__sm_103__tp4.layout``); this module changes only +what is stored. + +An expert projection of logical shape ``[N, K]`` is stored as four tensors:: + + .weight U8 [N, K/2] two E2M1 codes per byte; the low nibble holds + the even K index, the high nibble the next one + .weight_scale F8_E4M3 [N, K/16] one block scale per 16 consecutive K elements + .weight_scale_2 F32 [] the projection's global weight scale + .input_scale F32 [] the static activation scale s + +Its logical weight is ``W[n, k] = e2m1(code[n, k]) * e4m3(weight_scale[n, k // 16]) * +weight_scale_2``. An activation row ``x`` is quantized with ``g = 1 / s`` (per 16-element block +``b``: ``scale_b = e4m3(g * max|x_b| / 6)``, ``code = e2m1(x * g / scale_b)``), and a projection +is ``alpha * (codes x block scales) @ (codes x block scales)^T`` with ``alpha = s * +weight_scale_2``. + +Rank r owns intermediate ``[640 r, 640 r + 640)`` of every expert, on whole 16-element blocks:: + + gate_proj, up_proj [2560, 5120] rows [640 r, 640 r + 640) packed [640, 2560], scales [640, 320] + down_proj [5120, 2560] cols [640 r, 640 r + 640) packed [5120, 320] = byte columns + [320 r, 320 r + 320); scales + [5120, 40] = scale columns + [40 r, 40 r + 40) + +Every rank holds every projection's two FP32 scalars unchanged. Each sparse layer keeps each +projection kind of its 256 experts in banks (``layers..experts.``): packed codes +``[256, rows, K/2]``, block scales in the GEMM's 128x4 swizzled order ``[256, rows * K/16]``, and +the four scalars ``input_scale``, ``weight_scale_2``, ``g`` and ``alpha`` ``[256]`` FP32. +""" + +from __future__ import annotations + +import dataclasses +import fnmatch +import re +from collections import OrderedDict +from typing import Any, Mapping, Optional + +from ..checkpoint_layout_bf16__sm_103__tp4 import layout as bf16_layout +from ..checkpoint_layout_bf16__sm_103__tp4.layout import ( + BF16, + EMBED, + F32, + FINAL_NORM, + HEAD, + WORLD_SIZE, + TensorSpec, + check_config, + check_rank, +) + +U8 = "U8" +E4M3 = "F8_E4M3" +#: Elements per NVFP4 block scale. +GROUP = 16 +PROJECTIONS = ("gate_proj", "up_proj", "down_proj") +QUANT_ALGO = "NVFP4" +#: The quantizer's 128-row x 4-column scale tile. +SWIZZLE_ROWS, SWIZZLE_COLS = 128, 4 + +__all__ = [ + "BF16", + "E4M3", + "EMBED", + "F32", + "FINAL_NORM", + "GROUP", + "HEAD", + "PROJECTIONS", + "U8", + "WORLD_SIZE", + "QuantSpec", + "TensorSpec", + "check_config", + "check_partition", + "check_rank", + "check_recipe", + "ownership_ledger", + "stored_entries", + "swizzled_numel", + "tensor_specs", +] + + +def swizzled_numel(rows: int, cols: int) -> int: + """Bytes of a ``[rows, cols]`` block-scale matrix in the 128x4 swizzled layout (padded).""" + pad = lambda v, m: -(-v // m) * m # noqa: E731 + return pad(rows, SWIZZLE_ROWS) * pad(cols, SWIZZLE_COLS) + + +@dataclasses.dataclass(frozen=True) +class QuantSpec: + """One NVFP4 expert projection: logical geometry, four stored tensors, four-rank rule.""" + + name: str # "model.layers..mlp.experts.." + shape: tuple # logical [N, K] + layer: int + expert: int + proj: str # one of PROJECTIONS + split_dim: int # 0: output rows (gate/up), 1: input columns (down) + + @property + def family(self) -> str: + return "expert_" + self.proj.removesuffix("_proj") + + @property + def bank(self) -> str: + return f"layers.{self.layer}.experts.{self.proj.removesuffix('_proj')}" + + @property + def packed_shape(self) -> tuple: + return (self.shape[0], self.shape[1] // 2) + + @property + def scale_shape(self) -> tuple: + return (self.shape[0], self.shape[1] // GROUP) + + def stored(self) -> "OrderedDict[str, tuple]": + """The four stored tensors: name -> (dtype, shape).""" + return OrderedDict( + ( + (self.name + ".weight", (U8, self.packed_shape)), + (self.name + ".weight_scale", (E4M3, self.scale_shape)), + (self.name + ".weight_scale_2", (F32, ())), + (self.name + ".input_scale", (F32, ())), + ) + ) + + def interval(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + """``(dim, start, stop)`` of rank ``rank``'s logical slice, on whole quantization blocks.""" + check_rank(rank, world_size) + extent = self.shape[self.split_dim] + if extent % (world_size * GROUP): + raise ValueError( + f"{self.name}: dim {self.split_dim} of {extent} does not split into " + f"{world_size} whole {GROUP}-element blocks" + ) + size = extent // world_size + return self.split_dim, rank * size, (rank + 1) * size + + def packed_interval(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + """``(dim, start, stop)`` of the rank's slice of the packed U8 tensor (byte columns).""" + dim, start, stop = self.interval(rank, world_size) + return (dim, start, stop) if dim == 0 else (dim, start // 2, stop // 2) + + def scale_interval(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + """``(dim, start, stop)`` of the rank's slice of the block-scale tensor (scale columns).""" + dim, start, stop = self.interval(rank, world_size) + return (dim, start, stop) if dim == 0 else (dim, start // GROUP, stop // GROUP) + + def local_shape(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + dim, start, stop = self.interval(rank, world_size) + shape = list(self.shape) + shape[dim] = stop - start + return tuple(shape) + + def local_packed_shape(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + rows, cols = self.local_shape(rank, world_size) + return (rows, cols // 2) + + def local_scale_shape(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + rows, cols = self.local_shape(rank, world_size) + return (rows, cols // GROUP) + + +def tensor_specs(config: Any, layers: Optional[tuple] = None) -> "OrderedDict[str, Any]": + """Every logical tensor of the selected layers and the outer tensors. + + Keys are checkpoint names for higher-precision tensors (``TensorSpec``, stored BF16 or F32) + and projection prefixes for NVFP4 expert projections (``QuantSpec``). Geometry and + intervals are the BF16 layout's; the stored dtype of every non-expert matrix is BF16. + """ + specs: "OrderedDict[str, Any]" = OrderedDict() + for name, spec in bf16_layout.tensor_specs(config, layers).items(): + if spec.expert is not None: + prefix = name.removesuffix(".weight") + proj = prefix.rsplit(".", 1)[1] + specs[prefix] = QuantSpec( + prefix, tuple(spec.shape), spec.layer, spec.expert, proj, spec.split_dim + ) + elif spec.stored_dtype == bf16_layout.FP8: + specs[name] = dataclasses.replace(spec, stored_dtype=BF16) + else: + specs[name] = spec + return specs + + +def stored_entries(specs: Mapping[str, Any]) -> "OrderedDict[str, tuple]": + """Every stored tensor ``specs`` needs: name -> (dtype, shape).""" + out: "OrderedDict[str, tuple]" = OrderedDict() + for spec in specs.values(): + if isinstance(spec, QuantSpec): + out.update(spec.stored()) + else: + out[spec.name] = (spec.stored_dtype, tuple(spec.shape)) + return out + + +def check_partition(specs: Mapping[str, Any], world_size: int = WORLD_SIZE) -> None: + """Raise unless every split tensor's rank intervals tile its axis exactly, in rank order. + + For an expert projection the logical, packed-byte and block-scale intervals must each tile + their axis, and every logical boundary must fall on a quantization block boundary. + """ + bf16_layout.check_partition( + {n: s for n, s in specs.items() if not isinstance(s, QuantSpec)}, world_size + ) + for spec in specs.values(): + if not isinstance(spec, QuantSpec): + continue + for kind, extent, interval in ( + ("logical", spec.shape[spec.split_dim], spec.interval), + ("packed", spec.packed_shape[spec.split_dim], spec.packed_interval), + ("scale", spec.scale_shape[spec.split_dim], spec.scale_interval), + ): + cursor = 0 + for rank in range(world_size): + _, start, stop = interval(rank, world_size) + if start != cursor or stop <= start: + raise ValueError( + f"{spec.name}: rank {rank} {kind} interval [{start}, {stop}) leaves a gap " + "or overlap" + ) + cursor = stop + if cursor != extent: + raise ValueError(f"{spec.name}: {kind} intervals cover {cursor} of {extent}") + + +def ownership_ledger(specs: Mapping[str, Any], world_size: int = WORLD_SIZE) -> list: + """One plain-data row per logical tensor: its rule and every rank's intervals and storage. + + Higher-precision rows are the BF16 layout's. An expert row records, per rank, the logical, + packed-byte and block-scale intervals, the bank and the slot (the global expert id). + """ + rows = bf16_layout.ownership_ledger( + OrderedDict((n, s) for n, s in specs.items() if not isinstance(s, QuantSpec)), world_size + ) + for spec in specs.values(): + if not isinstance(spec, QuantSpec): + continue + ranks = [] + for rank in range(world_size): + dim, start, stop = spec.interval(rank, world_size) + _, pstart, pstop = spec.packed_interval(rank, world_size) + _, sstart, sstop = spec.scale_interval(rank, world_size) + ranks.append( + { + "rank": rank, + "replicated": False, + "dim": dim, + "start": start, + "stop": stop, + "packed": [pstart, pstop], + "scale": [sstart, sstop], + "local_packed_shape": list(spec.local_packed_shape(rank, world_size)), + "local_scale_shape": list(spec.local_scale_shape(rank, world_size)), + "scalars": "replicated", + "bank": spec.bank, + "slot": spec.expert, + } + ) + rows.append( + { + "name": spec.name, + "family": spec.family, + "layer": spec.layer, + "expert": spec.expert, + "shape": list(spec.shape), + "stored": {k: [d, list(s)] for k, (d, s) in spec.stored().items()}, + "ranks": ranks, + } + ) + return rows + + +def _module_names(specs: Mapping[str, Any]) -> tuple: + """(quantized expert modules, higher-precision Linear modules) of ``specs``.""" + quantized, kept = [], [] + for name, spec in specs.items(): + if isinstance(spec, QuantSpec): + quantized.append(name) + elif name.endswith(".weight") and len(spec.shape) == 2 and name != EMBED: + kept.append(name.removesuffix(".weight")) + return quantized, kept + + +def check_recipe(quant_config: Any, specs: Mapping[str, Any]) -> dict: + """Raise ValueError unless ``hf_quant_config.json`` is the recipe this layout computes. + + ModelOpt NVFP4, group size 16, no KV-cache quantization, and an exclusion list (``fnmatch`` + patterns over module names) that excludes every higher-precision Linear module of ``specs`` + (attention, router, dense feed-forward, head) and no expert projection. Returns a summary. + """ + if not isinstance(quant_config, Mapping): + raise ValueError("hf_quant_config.json must be a JSON object") + q = quant_config.get("quantization") + if not isinstance(q, Mapping): + raise ValueError("hf_quant_config.json has no 'quantization' object") + if q.get("quant_algo") != QUANT_ALGO: + raise ValueError(f"unsupported quant_algo {q.get('quant_algo')!r}; expected {QUANT_ALGO}") + if q.get("group_size") != GROUP or type(q.get("group_size")) is not int: + raise ValueError(f"unsupported group_size {q.get('group_size')!r}; expected {GROUP}") + if q.get("kv_cache_quant_algo") is not None: + raise ValueError( + f"unsupported kv_cache_quant_algo {q.get('kv_cache_quant_algo')!r}; expected none" + ) + patterns = q.get("exclude_modules") + if not isinstance(patterns, list) or not all(isinstance(p, str) for p in patterns): + raise ValueError("exclude_modules must be a list of strings") + excluded = re.compile("|".join(f"(?:{fnmatch.translate(p)})" for p in patterns) or "(?!)") + quantized, kept = _module_names(specs) + wrongly_excluded = [m for m in quantized if excluded.match(m)] + if wrongly_excluded: + raise ValueError(f"exclude_modules excludes quantized experts: {wrongly_excluded[:4]}") + not_excluded = [m for m in kept if not excluded.match(m)] + if not_excluded: + raise ValueError(f"higher-precision modules are not excluded: {not_excluded[:4]}") + return { + "quant_algo": QUANT_ALGO, + "group_size": GROUP, + "kv_cache_quant_algo": None, + "exclude_patterns": len(patterns), + "quantized_modules": len(quantized), + "excluded_linear_modules": len(kept), + } diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/loader.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/loader.py new file mode 100644 index 000000000000..545171d7c91d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/loader.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Validation and bounded, byte-preserving reads of an NVFP4-expert safetensors checkpoint. + +The index/header catalog and the two-buffer tile reader are the BF16 layout's +(``SafetensorsCatalog``, ``TileReader``): metadata only at construction, at most two pinned +staging buffers of ``TILE_BYTES``, never a whole tensor, shard or model on the host. + +``validate`` checks the catalog against every stored tensor the layout needs before anything +is loaded: present, nothing else present, exact dtype and shape, and a byte span matching them. + +``read_packed_into`` copies the rank's rows (and, for a column split, its byte columns) of a +stored 2-D byte tensor -- packed U8 codes or E4M3 block-scale bytes -- into device storage +unchanged. A column split reads whole rows tile by tile and keeps its columns on the device. +Values are never decoded, rescaled or cast here. ``read_scalars`` reads FP32 scalars with one +bounded read per contiguous run. +""" + +from __future__ import annotations + +import hashlib +import math +import struct +from typing import Mapping, Optional + +import torch + +from ..checkpoint_layout_bf16__sm_103__tp4.loader import ( + TILE_BYTES, + CheckpointFormatError, + SafetensorsCatalog, + TileReader, + read_slice_into, +) +from .layout import BF16, E4M3, F32, U8 + +__all__ = [ + "TILE_BYTES", + "CheckpointFormatError", + "SafetensorsCatalog", + "TileReader", + "check_block_scales", + "read_packed_into", + "read_scalars", + "read_slice_into", + "validate", +] + +_ITEMSIZE = {U8: 1, E4M3: 1, BF16: 2, F32: 4} +#: Scalar runs read in one request when they are at most this many bytes apart end to end. +_SCALAR_RUN_BYTES = 1 << 20 + + +class _Rows: + """A stored tensor seen as ``rows`` rows of ``row_bytes`` bytes, for ``TileReader.tiles``.""" + + __slots__ = ("name", "path", "begin", "end", "row_bytes") + + def __init__(self, name: str, path: str, begin: int, end: int, row_bytes: int): + self.name, self.path, self.begin, self.end, self.row_bytes = ( + name, + path, + begin, + end, + row_bytes, + ) + + +def validate(catalog: SafetensorsCatalog, expected: Mapping[str, tuple], *, complete=True) -> None: + """Raise CheckpointFormatError unless the stored tensors are exactly ``expected``. + + ``expected`` maps each stored name to ``(dtype, shape)`` (``layout.stored_entries``). With + ``complete`` any other stored tensor is rejected. + """ + missing = sorted(set(expected) - set(catalog.entries)) + if missing: + raise CheckpointFormatError(f"{len(missing)} expected tensors are missing: {missing[:8]}") + if complete: + extra = sorted(set(catalog.entries) - set(expected)) + if extra: + raise CheckpointFormatError(f"{len(extra)} unsupported tensors: {extra[:8]}") + for name, (dtype, shape) in expected.items(): + entry = catalog.entries[name] + if entry.dtype != dtype or tuple(entry.shape) != tuple(shape): + raise CheckpointFormatError( + f"{name} is stored as {entry.dtype}{list(entry.shape)}; expected {dtype}{list(shape)}" + ) + if entry.end - entry.begin != math.prod(shape) * _ITEMSIZE[dtype]: + raise CheckpointFormatError(f"{name}: byte span does not match {dtype}{list(shape)}") + + +def _sync(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.current_stream(device).synchronize() + + +@torch.no_grad() +def read_packed_into( + catalog: SafetensorsCatalog, + reader: TileReader, + name: str, + rows: tuple, + cols: Optional[tuple], + dst: torch.Tensor, + *, + digest: Optional["hashlib._Hash"] = None, +) -> None: + """Copy rows ``[rows[0], rows[1])`` (byte columns ``cols``, or all) of 2-D byte tensor ``name``. + + ``dst`` is uint8 ``[rows, cols]`` storage (any device), filled with the stored bytes + unchanged. ``digest`` (optional) is updated with the source bytes read, in tile order. + """ + entry = catalog.entries[name] + if entry.dtype not in (U8, E4M3) or len(entry.shape) != 2: + raise CheckpointFormatError(f"{name}: expected a 2-D U8/F8_E4M3 tensor, got {entry.dtype}") + row_bytes = entry.shape[1] + r0, r1 = rows + c0, c1 = (0, row_bytes) if cols is None else cols + if not (0 <= r0 < r1 <= entry.shape[0] and 0 <= c0 < c1 <= row_bytes): + raise ValueError(f"{name}: slice rows {rows} cols {cols} outside {list(entry.shape)}") + if dst.dtype != torch.uint8 or tuple(dst.shape) != (r1 - r0, c1 - c0): + raise ValueError( + f"{name}: destination {dst.dtype}{list(dst.shape)} != uint8{[r1 - r0, c1 - c0]}" + ) + view = _Rows(entry.name, entry.path, entry.begin, entry.end, row_bytes) + for t0, t1, raw in reader.tiles(view, r0, r1): + tile = raw.to(dst.device, non_blocking=True).view(t1 - t0, row_bytes) + dst[t0 - r0 : t1 - r0].copy_(tile[:, c0:c1]) + _sync(dst.device) + if digest is not None: + digest.update(f"{entry.name}:{t0}:{t1}|".encode()) + digest.update(memoryview(raw.numpy())) + + +def read_scalars(catalog: SafetensorsCatalog, reader: TileReader, names) -> dict: + """FP32 scalar tensors ``names`` -> ``{name: float}``, one bounded read per contiguous run.""" + entries = sorted((catalog.entries[n] for n in names), key=lambda e: (e.path, e.begin)) + for e in entries: + if e.dtype != F32 or tuple(e.shape) != () or e.end - e.begin != 4: + raise CheckpointFormatError(f"{e.name}: expected an F32 scalar, got {e.dtype}{e.shape}") + out = {} + limit = min(_SCALAR_RUN_BYTES, reader.tile_bytes) + i = 0 + while i < len(entries): + j = i + 1 + while ( + j < len(entries) + and entries[j].path == entries[i].path + and entries[j].end - entries[i].begin <= limit + ): + j += 1 + run = entries[i:j] + begin, end = run[0].begin, max(e.end for e in run) + view = _Rows(run[0].name, run[0].path, begin, end, end - begin) + for _, _, raw in reader.tiles(view, 0, 1): + data = bytes(raw.numpy()) + for e in run: + off = e.begin - begin + out[e.name] = struct.unpack(" None: + """Raise CheckpointFormatError on a NaN or negative E4M3 block scale (zero is allowed).""" + magnitude = scale_bytes & 0x7F + bad_nan = magnitude == 0x7F + bad_negative = ((scale_bytes & 0x80) != 0) & (magnitude != 0) + if bool((bad_nan | bad_negative).any()): + raise CheckpointFormatError(f"{name}: block scales hold NaN or negative E4M3 values") diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/model.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/model.py new file mode 100644 index 000000000000..5ee2a9e3ba1a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/model.py @@ -0,0 +1,498 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam NVFP4-expert checkpoint layout / sm_103 / tp4: packed experts and representative blocks. + +A component, not a routed target. There is no ``routing.py`` and no ``ForCausalLM``: a caller on +each of four ranks constructs this class with the checkpoint's configuration and its rank's +``Mapping`` (attention TP4, MoE TP4/EP1, PP1), loads a caller-supplied ModelOpt NVFP4 checkpoint +directory into the rank's slices, and runs representative computations. It does not run a +full-depth forward, prefill/decode loop or sampler. + +This is the BF16 layout component (``checkpoint_layout_bf16__sm_103__tp4``) with NVFP4 expert +storage and execution. Everything outside the routed experts -- ownership, the embedding, +attention on TRTLLM attention and KVCacheManagerV2 metadata, the dense feed-forward, the +replicated FP32 router with its all-rank agreement check, dispatch/combine, the norms, depth +residuals, head and every collective -- is that component's code, unchanged, on the +checkpoint's BF16/FP32 tensors copied byte for byte. The experts stay packed: + + storage per sparse layer and projection kind, banks indexed by global expert id: packed + E2M1 codes ``[256, rows, K/2]`` u8 and E4M3 block scales in the GEMM's 128x4 + swizzled order ``[256, rows * K/16]`` u8 (the rank's slice, see ``layout``), and + ``input_scale``, ``weight_scale_2``, ``g = 1/input_scale``, ``alpha = + input_scale * weight_scale_2`` ``[256]`` fp32 + compute ``expert.nvfp4_swiglu_expert``: per selected expert, activation quantization, + three CUTLASS NVFP4 GEMMs and the fp32 SwiGLU, as the expert provider of + ``beam_dispatch_combine``; the combination (bf16 routes, fp32 ascending-expert + sum, x2.5, bf16) is unchanged, and the one MoE all-reduce sums the four bf16 rank + partials in fp32 and rounds the sum to bf16 once + +Loading streams bounded tiles straight into that storage and is a transaction: metadata +(config.json, hf_quant_config.json, every index/header entry) is validated against the whole +model before any storage changes, and the component is usable only after every rank loaded +every selected tensor. Every parameter, bank and derived scalar is a registered parameter, so +``weight_state`` covers packed codes, block scales and scalars. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import struct +from collections import OrderedDict +from typing import Any, Mapping, Optional, Sequence + +import torch +from torch import nn + +from tensorrt_llm._torch.distributed import AllReduce +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.mapping import Mapping as ParallelMapping + +from ..checkpoint_layout_bf16__sm_103__tp4.model import ( + MAX_TOKENS, + BeamCheckpointLayoutBf16Sm103Tp4, + LoadError, + _attr, + _check_mapping, + _rec, + check_max_tokens, +) +from ..moe_dispatch_combine_bf16__sm_103__tp1 import dispatch_combine +from . import expert as nvfp4_expert +from .layout import ( + EMBED, + PROJECTIONS, + WORLD_SIZE, + QuantSpec, + check_config, + check_partition, + check_recipe, + ownership_ledger, + stored_entries, + swizzled_numel, + tensor_specs, +) +from .loader import ( + CheckpointFormatError, + SafetensorsCatalog, + TileReader, + check_block_scales, + read_packed_into, + read_scalars, + read_slice_into, + validate, +) + +_SM = (10, 3) +QUANT_CONFIG_FILE = "hf_quant_config.json" +#: Bank fields of one projection kind of one sparse layer. +BANK_FIELDS = ("weight", "weight_scale", "input_scale", "weight_scale_2", "g", "alpha") +_KINDS = tuple(p.removesuffix("_proj") for p in PROJECTIONS) + + +def _f32_bits(value: float) -> int: + return struct.unpack(" torch.Tensor: + """The rank-local storage of higher-precision tensor ``name``.""" + spec = self.specs[name] + if isinstance(spec, QuantSpec): + raise ValueError(f"{name} is an NVFP4 expert projection; see expert_storage()") + return getattr(self, _attr(name)) + + def bank(self, bank: str, field: str = "weight") -> torch.Tensor: + """The current ``[256, ...]`` storage of one field of one sparse layer's expert bank.""" + if field not in BANK_FIELDS: + raise ValueError(f"unknown bank field {field!r}; expected one of {BANK_FIELDS}") + return getattr(self, _attr(f"{bank}.{field}")) + + def expert_bank(self, layer: int) -> dict: + """``{"gate"|"up"|"down": {field: [256, ...] storage}}`` of sparse layer ``layer``.""" + self._check_usable(layer) + if not self.schedule.is_sparse(layer): + raise ValueError(f"layer {layer} has a dense feed-forward") + return { + kind: {f: self.bank(f"layers.{layer}.experts.{kind}", f) for f in BANK_FIELDS} + for kind in _KINDS + } + + def expert_storage(self, layer: int, expert: int, proj: str) -> dict: + """Views of one expert projection's rank-local storage (``proj`` in ``PROJECTIONS``).""" + spec = self.specs.get(f"model.layers.{layer}.mlp.experts.{expert}.{proj}") + if not isinstance(spec, QuantSpec): + raise ValueError(f"layer {layer} expert {expert} {proj!r} is not held") + return {f: self.bank(spec.bank, f)[expert] for f in BANK_FIELDS} + + def ledger(self) -> list: + """This component's ownership rows (see ``layout.ownership_ledger``).""" + return ownership_ledger(self.specs) + + # ------------------------------------------------------------------ loading + + def _load_expert(self, catalog, reader, spec: QuantSpec, digest) -> dict: + dim, p0, p1 = spec.packed_interval(self.rank) + _, s0, s1 = spec.scale_interval(self.rank) + if dim == 0: + rows, packed_cols, scale_cols = (p0, p1), None, None + else: + rows, packed_cols, scale_cols = (0, spec.shape[0]), (p0, p1), (s0, s1) + packed = self.bank(spec.bank, "weight")[spec.expert] + read_packed_into( + catalog, reader, spec.name + ".weight", rows, packed_cols, packed, digest=digest + ) + linear = torch.empty( + spec.local_scale_shape(self.rank), dtype=torch.uint8, device=self.device + ) + read_packed_into( + catalog, reader, spec.name + ".weight_scale", rows, scale_cols, linear, digest=digest + ) + check_block_scales(spec.name + ".weight_scale", linear) + swizzled = torch.ops.trtllm.block_scale_interleave(linear).reshape(-1) + dst = self.bank(spec.bank, "weight_scale")[spec.expert] + if swizzled.numel() != dst.numel(): + raise CheckpointFormatError( + f"{spec.name}: swizzled scales hold {swizzled.numel()} bytes, storage {dst.numel()}" + ) + dst.copy_(swizzled) + torch.cuda.current_stream(self.device).synchronize() + return { + "interval": list(spec.interval(self.rank)), + "packed_interval": [dim, p0, p1], + "scale_interval": [dim, s0, s1], + "local_packed_shape": list(packed.shape), + "local_scale_shape": list(linear.shape), + "bank": spec.bank, + "slot": spec.expert, + "source_sha256": digest.hexdigest() if digest is not None else None, + } + + def _load_scalars(self, catalog, reader, layer: int) -> dict: + experts = [s for s in self.specs.values() if isinstance(s, QuantSpec) and s.layer == layer] + names = [s.name + suffix for s in experts for suffix in (".input_scale", ".weight_scale_2")] + values = read_scalars(catalog, reader, names) + bad = [n for n, v in values.items() if not (v > 0.0 and v != float("inf"))] + if bad: + raise CheckpointFormatError(f"non-finite or non-positive scalars: {bad[:4]}") + n_experts = self.config["num_experts"] + shared = [False] * n_experts + for kind in _KINDS: + s = torch.empty(n_experts, dtype=torch.float32) + w2 = torch.empty(n_experts, dtype=torch.float32) + for e in range(n_experts): + p = f"model.layers.{layer}.mlp.experts.{e}.{kind}_proj" + s[e], w2[e] = values[p + ".input_scale"], values[p + ".weight_scale_2"] + g = torch.ones(n_experts, dtype=torch.float32) / s + alpha = s * w2 + bank = f"layers.{layer}.experts.{kind}" + for field, value in ( + ("input_scale", s), + ("weight_scale_2", w2), + ("g", g), + ("alpha", alpha), + ): + self.bank(bank, field).copy_(value) + for e in range(n_experts): + p = f"model.layers.{layer}.mlp.experts.{e}." + shared[e] = _f32_bits(values[p + "gate_proj.input_scale"]) == _f32_bits( + values[p + "up_proj.input_scale"] + ) + self._shared_fc1[layer] = shared + return {"scalars": len(values), "gate_up_input_scales_shared": sum(shared)} + + @torch.no_grad() + def load(self, checkpoint_dir: str, *, record_digests: bool = True) -> dict: + """Stream the selected tensors of a caller-supplied checkpoint into rank-local storage. + + Collective: every rank calls it. config.json must equal the constructed configuration, + hf_quant_config.json must be the NVFP4 recipe (``layout.check_recipe``), and every + stored tensor of the whole model must be exactly what the layout needs before any + storage changes. Expert codes and block-scale bytes are copied unchanged (block scales + then swizzled), scalars must be finite and positive, block scales finite and + non-negative. The component is unusable from the first write until every rank reports a + complete load; a failure on any rank raises ``LoadError`` on all ranks. Returns the + load record (recipe, coverage, source digests, bytes read). + """ + error: Optional[BaseException] = None + catalog = None + recipe = None + try: + with open(os.path.join(checkpoint_dir, "config.json")) as f: + stored = json.load(f) + if stored != self.config: + differing = sorted( + k for k in set(stored) | set(self.config) if stored.get(k) != self.config.get(k) + ) + raise CheckpointFormatError( + f"config.json differs from the constructed configuration: {differing[:8]}" + ) + with open(os.path.join(checkpoint_dir, QUANT_CONFIG_FILE)) as f: + quant = json.load(f) + full = tensor_specs(self.config) + recipe = check_recipe(quant, full) + catalog = SafetensorsCatalog(checkpoint_dir) + validate(catalog, stored_entries(full), complete=True) + except Exception as exc: # Keep rank-local failures inside the collective protocol. + error = exc + if not self._agree(error is None): + raise LoadError( + f"rank {self.rank}: checkpoint rejected before loading: {error!r}" + ) from error + + self.usable = False + self.load_record = None + self._shared_fc1 = {} + coverage = OrderedDict() + scalars = OrderedDict() + reader = None + bytes_read = reads = 0 + + def first_entry(spec): + return catalog.entries[ + spec.name + ".weight" if isinstance(spec, QuantSpec) else spec.name + ] + + try: + reader = TileReader() + order = sorted( + self.specs.values(), key=lambda s: (first_entry(s).path, first_entry(s).begin) + ) + sparse = [layer for layer in self.layers if self.schedule.is_sparse(layer)] + for spec in order: + digest = hashlib.sha256() if record_digests else None + if isinstance(spec, QuantSpec): + coverage[spec.name] = self._load_expert(catalog, reader, spec, digest) + continue + dst = self.tensor(spec.name) + read_slice_into(catalog, reader, spec, self.rank, dst, digest=digest) + coverage[spec.name] = { + "interval": spec.interval(self.rank), + "local_shape": list(dst.shape), + "dtype": str(dst.dtype), + "source_sha256": digest.hexdigest() if digest is not None else None, + } + for layer in sparse: + scalars[layer] = self._load_scalars(catalog, reader, layer) + except Exception as exc: # Every rank must reach the completion agreement. + error = exc + finally: + if reader is not None: + bytes_read, reads = reader.bytes_read, reader.reads + try: + reader.close() + except Exception as exc: # Cleanup failure must also reach the agreement. + if error is None: + error = exc + complete = ( + error is None and list(coverage) == [s.name for s in order] and list(scalars) == sparse + ) + if not self._agree(complete): + self._shared_fc1 = {} + raise LoadError( + f"rank {self.rank}: load did not complete on every rank: {error!r}" + ) from error + self.load_record = { + "rank": self.rank, + "checkpoint_dir": os.fspath(checkpoint_dir), + "layers": list(self.layers), + "recipe": recipe, + "contract": dict(nvfp4_expert.CONTRACT), + "tensors": len(coverage), + "expert_projections": len(self.expert_names), + "scalars": dict(scalars), + "bytes_read": bytes_read, + "reads": reads, + "coverage": coverage, + } + self.usable = True + return self.load_record + + # ------------------------------------------------------------------ computations + + @torch.no_grad() + def moe(self, layer: int, h: torch.Tensor, capture: Optional[dict] = None) -> torch.Tensor: + """Normalized input ``[T, hidden]`` -> MoE output after the one all-reduce. + + The replicated FP32 router, the all-rank agreement on its input, ids and weights and + the combination are the BF16 component's; each selected expert runs + ``expert.nvfp4_swiglu_expert`` on this rank's packed slice. The rank partials are + reduced once, in FP32, then rounded to BF16. + """ + routing = self.route(layer, h) + ids, weights = routing["ids"], routing["weights"] + digests = self._digests_agree((h, ids, weights)) + bank = self.expert_bank(layer) + shared = self._shared_fc1[layer] + calls = [] if capture is not None else None + + def provider(expert_id: int, rows: torch.Tensor) -> torch.Tensor: + return nvfp4_expert.nvfp4_swiglu_expert(rows, bank, expert_id, shared[expert_id], calls) + + local = dispatch_combine.beam_dispatch_combine(h, ids, weights, provider) + # The four BF16 rank partials are summed once in FP32 and the sum is rounded once. + out = self.all_reduce(local.float()).to(torch.bfloat16) + if capture is not None: + _rec(capture, _moe_capture(routing, calls, local, out)) + capture["moe.agreement_sha256"] = digests + return out + + +def _flat_with_offsets(parts: list) -> tuple: + offsets = [0] + for p in parts: + offsets.append(offsets[-1] + p.numel()) + return torch.cat([p.reshape(-1) for p in parts]), torch.tensor(offsets, dtype=torch.int64) + + +def _moe_capture(routing: dict, calls: list, local: torch.Tensor, out: torch.Tensor) -> dict: + """Routing intermediates and every expert call's operands, in call (ascending expert) order.""" + ids = routing["ids"] + tokens = [(ids == c["expert"]).any(-1).nonzero().reshape(-1) for c in calls] + fc1_sf, fc1_off = _flat_with_offsets([c["sf1"] for c in calls]) + fc2_sf, fc2_off = _flat_with_offsets([c["sf2"] for c in calls]) + items = { + "moe.logits": routing["logits"], + "moe.scores": routing["scores"], + "moe.selection": routing["selection_scores"], + "moe.group_scores": routing["group_scores"], + "moe.groups": routing["groups"], + "moe.masked": routing["masked_selection_scores"], + "moe.ids": ids, + "moe.weights": routing["weights"], + "moe.call_expert": torch.cat( + [torch.full((c["rows"].shape[0],), c["expert"], dtype=torch.int64) for c in calls] + ), + "moe.call_token": torch.cat(tokens), + "moe.call_rows": torch.tensor([c["rows"].shape[0] for c in calls], dtype=torch.int64), + "moe.call_shared_fc1": torch.tensor([c["fc1_up"] is None for c in calls]), + "moe.call_scalars": torch.stack([c["scalars"] for c in calls]), + "moe.expert_in": torch.cat([c["rows"] for c in calls]), + "moe.fc1_fp4": torch.cat([c["fc1"] for c in calls]), + "moe.fc1_sf_swizzled": fc1_sf, + "moe.fc1_sf_offsets": fc1_off, + "moe.gate": torch.cat([c["gate"] for c in calls]), + "moe.up": torch.cat([c["up"] for c in calls]), + "moe.silu": torch.cat([c["silu"] for c in calls]), + "moe.product": torch.cat([c["product"] for c in calls]), + "moe.act": torch.cat([c["act"] for c in calls]), + "moe.fc2_fp4": torch.cat([c["fc2"] for c in calls]), + "moe.fc2_sf_swizzled": fc2_sf, + "moe.fc2_sf_offsets": fc2_off, + "moe.tp.y": torch.cat([c["y"] for c in calls]), + "moe.tp.local_out": local, + "moe.output": out, + "mlp_out": out, + } + unshared = [c for c in calls if c["fc1_up"] is not None] + if unshared: + sf, off = _flat_with_offsets([c["sf1_up"] for c in unshared]) + items.update( + { + "moe.fc1_up_expert": torch.tensor([c["expert"] for c in unshared]), + "moe.fc1_up_fp4": torch.cat([c["fc1_up"] for c in unshared]), + "moe.fc1_up_sf_swizzled": sf, + "moe.fc1_up_sf_offsets": off, + } + ) + return items + + +__all__ = [ + "BANK_FIELDS", + "BeamCheckpointLayoutNvfp4Sm103Tp4", + "LoadError", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/__init__.py new file mode 100644 index 000000000000..10b0e507d6c1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam decoder block 0 (dense, sliding attention) / sm_103 / tp1.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/block.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/block.py new file mode 100644 index 000000000000..3900e91aaca8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/block.py @@ -0,0 +1,293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam decoder blocks 0-2 / sm_103 / tp1: one dense, sliding-attention block, constructed directly. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine builds this class, and it serves no request. +A caller constructs it from the checkpoint's configuration for one of the dense blocks +0, 1 or 2, loads the block's fifteen tensors and runs one step of one BF16 request through +it -- a context step after any cached prefix, or a one-token generation step -- with the +attention batch state prepared on a ``KVCacheManagerV2``. The forward is +assembled from catalog entries, with the attention sublayer shared with the sparse block +component in ``../_attention.py``, and it does not call the built-in Beam modeling. + +The block, with block 0's residual coefficients (block i's are in ``depth_scales``):: + + h = rmsnorm(x, input_layernorm) + q, k, v, g = h @ q_proj^T, h @ k_proj^T, h @ v_proj^T, h @ gate_proj^T + q, k = rmsnorm(q, q_norm), rmsnorm(k, k_norm) per head, rounded to bf16 + q, k = rope(q, k) fp32 NeoX rotation, rounded + a = causal attention(q, k, v), window 512, scale 1/sqrt(128), 6 q heads per kv head + a = a * sigmoid(g + gate_bias) per head, bf16 steps + x = rmsnorm(a @ o_proj^T, post_attention_layernorm) / sqrt(2) + x / sqrt(2) + h = rmsnorm(x, pre_feedforward_layernorm) + m = (silu(h @ gate^T) * (h @ up^T)) @ down^T silu and product in fp32 + out = rmsnorm(m, post_feedforward_layernorm) / sqrt(3) + x * sqrt(2) / sqrt(3) + +Every RMSNorm computes its statistics, normalization and affine in fp32 and rounds once. +""" + +from __future__ import annotations + +import math +from collections import OrderedDict +from typing import Mapping + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501 + flashinfer_silu_and_mul, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.residual.beam_depth_residual import ( + beam_depth_residual, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + +from .. import _attention + +#: The configuration this component computes, as the checkpoint's config.json states it. +#: Anything else is rejected at construction rather than computed approximately. +SUPPORTED_CONFIG = { + "hidden_size": 5120, + "intermediate_size": 20480, + "num_attention_heads": 48, + "num_key_value_heads": 8, + "head_dim": 128, + "rms_norm_eps": 1e-6, + "hidden_act": "silu", + "attention_bias": False, + "attention_dropout": 0.0, + "sliding_window": 512, +} +SUPPORTED_LAYER_IDX = 0 +#: The dense blocks of the checkpoint's schedule; ``layer_idx`` fixes the residual coefficients. +SUPPORTED_LAYER_INDICES = (0, 1, 2) +SUPPORTED_ROPE = {"rope_type": "default", "rope_theta": 10000.0} + +#: The most tokens one call may append: the six-layer model's sequence capacity. Longer +#: calls are rejected; the request's whole KV length is bounded by its cache's capacity. +MAX_TOKENS = 1024 + +_SM = (10, 3) + + +def weight_shapes(config) -> "OrderedDict[str, tuple[int, ...]]": + """The block's fifteen tensors by checkpoint name, in checkpoint order, [out, in] for matrices.""" + hidden, inter = config.hidden_size, config.intermediate_size + q_width = config.num_attention_heads * config.head_dim + kv_width = config.num_key_value_heads * config.head_dim + return OrderedDict( + [ + ("input_layernorm.weight", (hidden,)), + ("self_attn.gate_bias", (config.num_attention_heads,)), + ("self_attn.q_proj.weight", (q_width, hidden)), + ("self_attn.k_proj.weight", (kv_width, hidden)), + ("self_attn.v_proj.weight", (kv_width, hidden)), + ("self_attn.gate_proj.weight", (q_width, hidden)), + ("self_attn.o_proj.weight", (hidden, q_width)), + ("self_attn.q_norm.weight", (config.head_dim,)), + ("self_attn.k_norm.weight", (config.head_dim,)), + ("pre_feedforward_layernorm.weight", (hidden,)), + ("mlp.gate_proj.weight", (inter, hidden)), + ("mlp.up_proj.weight", (inter, hidden)), + ("mlp.down_proj.weight", (hidden, inter)), + ("post_attention_layernorm.weight", (hidden,)), + ("post_feedforward_layernorm.weight", (hidden,)), + ] + ) + + +def _attribute(name: str) -> str: + return name.removesuffix(".weight").replace(".", "_") + + +def depth_scales(layer_idx: int, is_attention: bool) -> tuple[float, float]: + """(branch, residual) scales of a block's attention or feed-forward residual add. + + The n-th add of the residual stream scales the branch by 1 / sqrt(n) and the stream + by sqrt(n - 1) / sqrt(n); block i's attention add is the 2(i + 1)-th, its + feed-forward add the next one. + """ + num_adds = 2 * (layer_idx + 1) + (0 if is_attention else 1) + denom = math.sqrt(num_adds) + return 1.0 / denom, math.sqrt(num_adds - 1) / denom + + +def check_config(config, layer_idx: int) -> None: + """Raise ValueError naming the first field outside what this component computes.""" + if type(layer_idx) is not int or layer_idx not in SUPPORTED_LAYER_INDICES: + raise ValueError( + f"only dense blocks {SUPPORTED_LAYER_INDICES} are supported, got layer_idx={layer_idx!r}" + ) + for field, expected in SUPPORTED_CONFIG.items(): + actual = getattr(config, field, None) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this block computes {expected!r}") + layer_types = getattr(config, "layer_types", None) or [] + if len(layer_types) <= layer_idx or layer_types[layer_idx] != "sliding_attention": + raise ValueError( + f"block {layer_idx} must be a sliding_attention layer: " + f"{layer_types[layer_idx] if len(layer_types) > layer_idx else None!r}" + ) + mlp_types = getattr(config, "mlp_layer_types", None) or [] + if len(mlp_types) <= layer_idx or mlp_types[layer_idx] != "dense": + raise ValueError( + f"block {layer_idx} must have a dense MLP: " + f"{mlp_types[layer_idx] if len(mlp_types) > layer_idx else None!r}" + ) + rope = (getattr(config, "rope_parameters", None) or {}).get("sliding_attention") + if rope != SUPPORTED_ROPE: + raise ValueError( + f"unsupported sliding_attention rope parameters {rope!r}; need {SUPPORTED_ROPE!r}" + ) + + +class BeamDenseBlockBf16Sm103Tp1(nn.Module): + """Beam dense decoder block 0, 1 or 2: BF16 weights and activations, one GPU, one request step. + + ``forward(hidden_states [1, T, 5120] bf16, position_ids [1, T], attn_metadata)`` + returns ``[1, T, 5120]`` bf16 for 1 <= T <= ``MAX_TOKENS``, a prepared + ``TrtllmAttentionMetadata`` holding one step of one request on a BF16 + ``KVCacheManagerV2`` (8 kv heads of 128) that holds the block's logical layer + ``layer_idx`` -- a context step of T tokens after ``past`` cached ones (paged-context + FMHA when ``past > 0``) or a one-token generation step -- and positions + ``arange(past, past + T)``. Inputs outside that are rejected before any kernel runs; + batches of several requests and other blocks are not computed here. + """ + + def __init__( + self, config, layer_idx: int = SUPPORTED_LAYER_IDX, device: str | torch.device = "cuda" + ): + super().__init__() + check_config(config, layer_idx) + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this block is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.layer_idx = layer_idx + self.eps = config.rms_norm_eps + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.window = config.sliding_window + self.rope_theta = SUPPORTED_ROPE["rope_theta"] + self.attention_scales = depth_scales(layer_idx, is_attention=True) + self.feed_forward_scales = depth_scales(layer_idx, is_attention=False) + self._shapes = weight_shapes(config) + for name, shape in self._shapes.items(): + tensor = torch.empty(shape, dtype=torch.bfloat16, device=device) + self.register_parameter(_attribute(name), nn.Parameter(tensor, requires_grad=False)) + self._loaded = False + + def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": + """The block's tensors by checkpoint name, in checkpoint order.""" + return OrderedDict((name, getattr(self, _attribute(name))) for name in self._shapes) + + @torch.no_grad() + def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Copy exactly the block's fifteen tensors, by checkpoint name, BF16, [out, in].""" + if not isinstance(weights, Mapping): + raise TypeError("weights must be a mapping") + missing = [n for n in self._shapes if n not in weights] + unexpected = sorted(n for n in weights if n not in self._shapes) + if missing or unexpected: + raise ValueError(f"block weights mismatch: missing {missing}, unexpected {unexpected}") + for name, shape in self._shapes.items(): + tensor = weights[name] + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tuple(tensor.shape) != shape or tensor.dtype != torch.bfloat16: + raise ValueError( + f"{name}: got {tuple(tensor.shape)} {tensor.dtype}, need {shape} torch.bfloat16" + ) + if tensor.layout != torch.strided or tensor.device.type not in ("cpu", "cuda"): + raise ValueError(f"{name} must be a dense host or CUDA tensor") + self._loaded = False + for name in self._shapes: + getattr(self, _attribute(name)).copy_(weights[name]) + self._loaded = True + + def _check_step( + self, hidden_states: torch.Tensor, position_ids: torch.Tensor, md: TrtllmAttentionMetadata + ) -> None: + """Raise ValueError on anything this block would otherwise compute wrongly.""" + _attention.check_step( + hidden_states, + position_ids, + md, + layer_idx=self.layer_idx, + hidden_size=self.hidden_size, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + device=self.input_layernorm.device, + max_tokens=MAX_TOKENS, + ) + if not self._loaded: + raise ValueError("load_weights must complete before running the block") + + @property + def hidden_size(self) -> int: + return self._shapes["input_layernorm.weight"][0] + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """Project the prepared metadata onto thop_attention's explicit batch state.""" + return _attention.step_args(md, self.layer_idx) + + def attention(self, h: torch.Tensor, positions: torch.Tensor, step: dict) -> torch.Tensor: + """Normalized input [T, hidden] -> attention output [T, hidden], before post_attention_layernorm.""" + return _attention.attention( + h, + positions, + step, + q_proj=self.self_attn_q_proj, + k_proj=self.self_attn_k_proj, + v_proj=self.self_attn_v_proj, + gate_proj=self.self_attn_gate_proj, + o_proj=self.self_attn_o_proj, + q_norm=self.self_attn_q_norm, + k_norm=self.self_attn_k_norm, + gate_bias=self.self_attn_gate_bias, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + eps=self.eps, + rope_theta=self.rope_theta, + window=self.window, + ) + + def feed_forward(self, h: torch.Tensor) -> torch.Tensor: + """Normalized input [T, hidden] -> dense SwiGLU output [T, hidden], before its post-norm.""" + gate = cublas_mm(h, self.mlp_gate_proj.t()) + up = cublas_mm(h, self.mlp_up_proj.t()) + act = flashinfer_silu_and_mul(torch.cat((gate, up), dim=-1)) + return cublas_mm(act, self.mlp_down_proj.t()) + + @torch.no_grad() + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + attn_metadata: TrtllmAttentionMetadata, + ) -> torch.Tensor: + self._check_step(hidden_states, position_ids, attn_metadata) + step = self.step_args(attn_metadata) + num_tokens = hidden_states.shape[1] + x = hidden_states.view(num_tokens, self.hidden_size) + positions = position_ids.view(num_tokens) + + h = flashinfer_rmsnorm(x, self.input_layernorm, self.eps) + h = self.attention(h, positions, step) + h = flashinfer_rmsnorm(h, self.post_attention_layernorm, self.eps) + x = beam_depth_residual(h, x, *self.attention_scales) + + h = flashinfer_rmsnorm(x, self.pre_feedforward_layernorm, self.eps) + h = self.feed_forward(h) + h = flashinfer_rmsnorm(h, self.post_feedforward_layernorm, self.eps) + x = beam_depth_residual(h, x, *self.feed_forward_scales) + return x.view(1, num_tokens, self.hidden_size) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/model.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/model.py new file mode 100644 index 000000000000..c2eeb529fb83 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/model.py @@ -0,0 +1,352 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam full model, NVFP4 experts / sm_103 / tp4: token ids to logits through all 52 layers. + +The packed NVFP4 checkpoint-layout component (``checkpoint_layout_nvfp4__sm_103__tp4``) with a +full-depth forward. Every operation is the component's own: the masked embedding with its +all-reduce, each logical layer's ``block`` (attention on TRTLLM attention and the call's +``KVCacheManagerV2`` slot for that layer, the dense SwiGLU on layers 0-2, the replicated FP32 +router and packed W4A4 experts on layers 3-51, the four sandwich norms and the depth residuals), +the final norm, the untied head with its vocabulary all-gather and the FP32 softcap. This class +adds the order and the contract of one call:: + + x = embed(input_ids) replicated [T, 5120] bf16 + for layer in 0 .. 51: x = block(layer, x) in logical order, each exactly once + logits = head(final_norm(x[rows])) fp32 [len(rows), 151936], softcapped once + +One call is one engine step of one or more requests (``_attention.check_batch_metadata``): +the step's context requests first, then its generation requests, their new tokens +concatenated into the call's T rows. A context request appends its ``n`` tokens after the +``past`` tokens the cache already holds for it (none, a reused prefix or earlier chunks), a +generation request appends ``generation_rows`` tokens: one, or in a speculative +verification step ``P`` (its pending token and ``P - 1`` proposals). Each request's +``position_ids`` must be exactly ``past .. past + n - 1``. ``logit_rows`` selects the hidden +rows that reach the final norm and the head before either runs: ``None`` (default) is the +last row of each request, in request order -- the next-token logits a sampler needs -- +``"all"`` every row, ``"verification"`` the last row of each context request and every row +of each generation request (what a speculative verifier reads), or explicit row indices. +``block_output(layer, x)``, when given, receives each logical layer's completed block output +``x`` ``[T, hidden]`` (the residual stream after the block, before the next block or the +final norm) as soon as the block returns. + +The cache must be a BF16 ``KVCacheManagerV2`` holding all 52 logical layers in logical order +with 2 kv heads of 128 per rank, a 512-token window on every sliding layer and the whole +sequence (``max_seq_len``) on the full-attention layers 5, 11, 17, 23, 29, 35, 41, 47 and 51 +(``attention_windows``). Inputs, positions, metadata and cache are validated before any kernel +runs; a call with T above ``max_tokens`` is rejected (the default ``MAX_NUM_TOKENS`` matches +the supported deployment's native ``max_num_tokens``). + +Host synchronization: the position check reads one device comparison back per call (and the +``gather_ids`` check, when given, one more); the component's router agreement and expert +dispatch synchronize per sparse layer. +""" + +from __future__ import annotations + +from typing import Any, Callable, Mapping, Optional, Sequence, Union + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam import _attention +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_nvfp4__sm_103__tp4.model import ( # noqa: E501 + BeamCheckpointLayoutNvfp4Sm103Tp4, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm.bindings import DataType +from tensorrt_llm.mapping import Mapping as ParallelMapping + +#: Logical decoder layers, all executed by every call. +NUM_LAYERS = 52 +#: Tokens per call of the supported deployment (its native ``max_num_tokens``). +MAX_NUM_TOKENS = 8192 +#: KV cache page size of the supported deployment. +TOKENS_PER_BLOCK = 32 + +LogitRows = Union[None, str, Sequence[int]] + + +class BeamFullModelNvfp4Sm103Tp4(BeamCheckpointLayoutNvfp4Sm103Tp4): + """All 52 layers of the Beam model with packed NVFP4 experts on one of four TP ranks. + + ``config`` is the checkpoint's config.json (a mapping) and ``mapping`` this rank's + ``tensorrt_llm.mapping.Mapping`` (world size 4, TP4, PP1, MoE TP4/EP1, no attention DP); + every logical layer is held. ``load(checkpoint_dir)`` (collective) fills the weights; the + forward is usable only after every rank completed it. + """ + + def __init__( + self, + config: Mapping[str, Any], + mapping: ParallelMapping, + device: Optional[torch.device] = None, + *, + max_tokens: int = MAX_NUM_TOKENS, + ): + super().__init__(config, mapping, layers=None, device=device, max_tokens=max_tokens) + if self.layers != tuple(range(NUM_LAYERS)): + raise ValueError(f"the full model holds logical layers 0..{NUM_LAYERS - 1}") + self.vocab_size = config["vocab_size"] + + def attention_windows(self, max_seq_len: int) -> list: + """Per logical layer: the attention window its cache must keep (512 or ``max_seq_len``).""" + return [ + self.window if self.schedule.is_sliding(layer) else max_seq_len + for layer in range(NUM_LAYERS) + ] + + def check_cache_manager(self, manager) -> None: + """Raise ValueError unless ``manager`` holds every logical layer as this model needs. + + Required: BF16, ``head_dim`` 128, 2 kv heads per rank on every layer, logical layers + 0..51 in local slots 0..51, the sliding window on sliding layers and the whole sequence + on full-attention layers. Reads host-side manager state only. + """ + offsets = getattr(manager, "layer_offsets", None) + if not isinstance(offsets, Mapping) or dict(offsets) != { + layer: layer for layer in range(NUM_LAYERS) + }: + held = sorted(offsets) if isinstance(offsets, Mapping) else None + raise ValueError( + f"the KV cache manager must hold logical layers 0..{NUM_LAYERS - 1} in slots " + f"0..{NUM_LAYERS - 1}; it holds {held}" + ) + heads = list(getattr(manager, "num_kv_heads_per_layer", None) or []) + if ( + manager.dtype != DataType.BF16 + or manager.head_dim != self.head_dim + or heads != [self.local_kv_heads] * NUM_LAYERS + ): + raise ValueError( + f"the KV cache must be BF16 with {self.local_kv_heads} kv heads of " + f"{self.head_dim} per layer; got {manager.dtype}, head_dim {manager.head_dim}" + ) + windows = list(getattr(manager, "max_attention_window_vec", None) or []) + if len(windows) != NUM_LAYERS: + raise ValueError(f"the KV cache manager states {len(windows)} attention windows") + for layer, window in enumerate(windows): + if self.schedule.is_sliding(layer): + ok = window is not None and int(window) == self.window + else: + ok = window is None or int(window) >= manager.max_seq_len + if not ok: + kind = "sliding" if self.schedule.is_sliding(layer) else "full-attention" + raise ValueError( + f"logical layer {layer} is a {kind} layer; the cache keeps a window of " + f"{window} (max_seq_len {manager.max_seq_len})" + ) + + def _select_rows( + self, logit_rows: LogitRows, spans: Sequence[tuple], num_contexts: int = 0 + ) -> list: + num_tokens = sum(n for _, n in spans) + if logit_rows is None: + ends, end = [], 0 + for _, n in spans: + end += n + ends.append(end - 1) + return ends + if isinstance(logit_rows, str): + if logit_rows == "all": + return list(range(num_tokens)) + if logit_rows != "verification": + raise ValueError( + "logit_rows must be None, 'all', 'verification' or row indices, got " + f"{logit_rows!r}" + ) + rows, start = [], 0 + for i, (_, n) in enumerate(spans): + rows.extend([start + n - 1] if i < num_contexts else range(start, start + n)) + start += n + return rows + rows = list(logit_rows) + if ( + not rows + or any(type(r) is not int or not 0 <= r < num_tokens for r in rows) + or len(set(rows)) != len(rows) + ): + raise ValueError( + f"logit_rows must be distinct row indices in [0, {num_tokens}), got {logit_rows!r}" + ) + return rows + + def check_call( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attn_metadata: TrtllmAttentionMetadata, + *, + generation_rows: int = 1, + extra_kv_tokens: int = 0, + ) -> tuple: + """Validate one call; return the flat ids, the flat positions and the request spans. + + ``input_ids``: ``[T]`` or ``[1, T]`` torch.int64 ids in ``[0, vocab)`` on the model's + device, 1 <= T <= ``max_tokens``; ``attn_metadata``: one prepared step of its requests + on a manager that passes ``check_cache_manager``, ``generation_rows`` rows per + generation request and ``extra_kv_tokens`` reserved KV slots per request; + ``position_ids``: an integer tensor of T positions, each request's + ``past .. past + n - 1`` in request order. Returns the spans ``[(past, n), ...]`` of + ``_attention.check_batch_metadata``. Raises ValueError otherwise. + """ + if not isinstance(input_ids, torch.Tensor) or input_ids.dtype != torch.int64: + raise ValueError("input_ids must be a torch.int64 tensor") + if input_ids.dim() == 2 and input_ids.shape[0] == 1: + input_ids = input_ids[0] + if input_ids.dim() != 1 or not 1 <= input_ids.shape[0] <= self.max_tokens: + raise ValueError( + f"input_ids must be [T] or [1, T] with 1 <= T <= {self.max_tokens}, " + f"got {list(input_ids.shape)}" + ) + if input_ids.device != self.device: + raise ValueError("input_ids must be on the model's device") + num_tokens = input_ids.shape[0] + if ( + not isinstance(position_ids, torch.Tensor) + or position_ids.is_floating_point() + or position_ids.dtype == torch.bool + or position_ids.numel() != num_tokens + or position_ids.device != self.device + ): + raise ValueError( + f"position_ids must be an integer tensor of {num_tokens} positions on the " + "model's device" + ) + if not isinstance(attn_metadata, TrtllmAttentionMetadata): + raise ValueError( + f"attn_metadata must be TrtllmAttentionMetadata, got {type(attn_metadata).__name__}" + ) + if attn_metadata.kv_cache_manager is None: + raise ValueError("attn_metadata needs a KV cache manager") + self.check_cache_manager(attn_metadata.kv_cache_manager) + spans = _attention.check_batch_metadata( + attn_metadata, + num_tokens, + layer_idx=0, + num_kv_heads=self.local_kv_heads, + head_dim=self.head_dim, + generation_rows=generation_rows, + extra_kv_tokens=extra_kv_tokens, + ) + positions = position_ids.reshape(-1).to(torch.int64) + expected = torch.cat([torch.arange(past, past + n, dtype=torch.int64) for past, n in spans]) + if not torch.equal(positions, expected.to(self.device, non_blocking=True)): + got, start = positions.cpu(), 0 + for i, (past, n) in enumerate(spans): + if not torch.equal(got[start : start + n], expected[start : start + n]): + break + start += n + raise ValueError( + f"request {i} of the call appends tokens {past}..{past + n - 1} after {past} " + "cached tokens; its positions must be exactly those" + ) + return input_ids, positions, spans + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attn_metadata: TrtllmAttentionMetadata, + *, + logit_rows: LogitRows = None, + layer_capture: Optional[Callable[[int], Optional[dict]]] = None, + trace: Optional[list] = None, + generation_rows: int = 1, + extra_kv_tokens: int = 0, + block_output: Optional[Callable[[int, torch.Tensor], None]] = None, + gather_ids: Optional[torch.Tensor] = None, + return_hidden: bool = False, + ) -> Union[torch.Tensor, tuple]: + """One engine step through all 52 layers; FP32 softcapped logits ``[R, vocab]``. + + ``layer_capture(layer)`` may return a dict that receives that block's named + intermediates (``block``'s capture); ``trace``, when a list, receives one record per + executed stage in execution order (embedding with the step's request spans, each + logical layer with its kind and window, final norm, head). ``generation_rows`` and + ``extra_kv_tokens`` state the step's speculative shape (``check_call``); + ``block_output`` is called with each layer's completed block output. ``gather_ids``, + the engine's own logits rows when it states them, must name exactly the rows + ``logit_rows`` selects. With ``return_hidden`` the result is ``(logits, x)``, ``x`` + the last block's output ``[T, hidden]`` before the final norm. + """ + self._check_usable() + ids, positions, spans = self.check_call( + input_ids, + position_ids, + attn_metadata, + generation_rows=generation_rows, + extra_kv_tokens=extra_kv_tokens, + ) + num_tokens = ids.shape[0] + rows = self._select_rows(logit_rows, spans, attn_metadata.num_contexts) + if gather_ids is not None: + self._check_gather_ids(gather_ids, rows) + x = self.embed(ids) + if trace is not None: + trace.append({"stage": "embed", "tokens": num_tokens, "requests": spans}) + for layer in range(NUM_LAYERS): + capture = layer_capture(layer) if layer_capture is not None else None + x = self.block( + layer, + x, + positions, + attn_metadata, + capture, + generation_rows=generation_rows, + extra_kv_tokens=extra_kv_tokens, + ) + if block_output is not None: + block_output(layer, x) + if trace is not None: + sliding = self.schedule.is_sliding(layer) + trace.append( + { + "stage": "layer", + "layer": layer, + "kv_slot": _attention.local_layer_index(attn_metadata, layer), + "mlp": "sparse" if self.schedule.is_sparse(layer) else "dense", + "attention": "sliding" if sliding else "full", + "window": self.window if sliding else attn_metadata.max_seq_len, + "output_sq_norm": x.float().square().sum(), + } + ) + selected = x if len(rows) == num_tokens and rows == list(range(num_tokens)) else x[rows] + normed = self.final_norm(selected.contiguous()) + logits = self.head(normed) + if trace is not None: + trace.append({"stage": "final_norm", "rows": rows}) + trace.append({"stage": "head", "rows": len(rows), "vocab": int(logits.shape[-1])}) + return (logits, x) if return_hidden else logits + + def _check_gather_ids(self, gather_ids: torch.Tensor, rows: list) -> None: + """Raise ValueError unless the engine's ``gather_ids`` are exactly ``rows``, in order. + + One device comparison read back per call. + """ + if ( + not isinstance(gather_ids, torch.Tensor) + or gather_ids.is_floating_point() + or gather_ids.dim() != 1 + or gather_ids.numel() != len(rows) + or gather_ids.device != self.device + ): + raise ValueError( + f"the engine gathers logits rows " + f"{None if not isinstance(gather_ids, torch.Tensor) else list(gather_ids.shape)} " + f"on {getattr(gather_ids, 'device', None)}; this step's verifier reads " + f"{len(rows)} rows on {self.device}" + ) + expected = torch.tensor(rows, dtype=torch.int64).to(self.device, non_blocking=True) + if not torch.equal(gather_ids.to(torch.int64), expected): + raise ValueError( + f"the engine gathers logits rows {gather_ids.tolist()}; the step's requests " + f"need the rows {rows} (each context's last row, every generation row)" + ) + + +__all__ = [ + "MAX_NUM_TOKENS", + "NUM_LAYERS", + "TOKENS_PER_BLOCK", + "BeamFullModelNvfp4Sm103Tp4", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/modeling.py new file mode 100644 index 000000000000..4d273c5583aa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/modeling.py @@ -0,0 +1,451 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Routed ModelingV2 target: the Beam NVFP4 full model on the native PyTorch runtime. + +``ModelingV2BeamFullModelNvfp4Sm103Tp4`` is what ``AutoModelForCausalLM`` builds for a +``BeamForCausalLM`` NVFP4 checkpoint on four sm_103 ranks (attention TP4, MoE TP4/EP1, PP1, no +attention DP), without speculative decoding or with a DFlash 2 drafter, when ``modeling_v2`` +routes it. It owns exactly one target parameter set: the full-model component +(``model.BeamFullModelNvfp4Sm103Tp4``), constructed from the checkpoint's own ``config.json`` +and filled by ``weights.BeamNvfp4CheckpointLoader`` through the component's transactional +``load``; there is no second embedding, head or expert copy. + +Each engine step is the native runtime's batch of one or more requests (up to the engine's +``max_batch_size``), context requests first, dispatched once on the phase: + +* ``PrefillTarget``: a step with at least one context request (its prompt tokens, after any + reused prefix), possibly followed by generation requests; the logits of each request's + last row, or of every row when the engine asks for context logits; +* ``DecodeTarget``: a step of generation requests, one token each; their next-token logits. + +Both return FP32 softcapped logits ``[rows, 151936]`` in request order; the engine's +TorchSampler selects tokens from them. The KV cache is the engine's ``KVCacheManagerV2``; the +component validates its layers, windows and the step's metadata before any kernel runs. Eager +execution only: the expert provider synchronizes with the host, so CUDA graphs and the overlap +scheduler are not supported. + +With a DFlash speculative config (``routing.route`` admits only the DFlash 2 drafter on the +TRTLLM draft-attention backend) the target composes the native one-model pieces rather than a +second scheduler, verifier or cache: the drafter ``get_draft_model`` builds from the drafter +checkpoint (``modeling_dflash.DFlashForCausalLM``), the ``DFlashWorker`` ``get_spec_worker`` +builds, and the engine's separate draft ``KVCacheManagerV2`` pool. Every step is then a +``VerifyTarget`` step: one target forward over the context rows and each generation request's +``P = runtime_draft_len + 1`` rows (the pending token and its proposals), with the completed +block outputs of the drafter's capture layers copied into the step's DFlash capture buffer and +the FP32 softcapped logits of the verification rows (each context's last row, every generation +row) handed to the worker, which accepts the proposals equal to the target's argmax, drafts the +next block and returns the native result dictionary. The drafter's logits processor is the +checkpoint's ``final_logit_softcapping``, applied once (``DFlashSoftcapLogitsProcessor``); its +embedding and vocabulary-parallel head are views of the target's own TP vocabulary slices. +""" + +from __future__ import annotations + +import json +import os +import types +from typing import Optional + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_nvfp4__sm_103__tp4.layout import ( # noqa: E501 + EMBED, + HEAD, +) +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.logger import logger + +from .model import MAX_NUM_TOKENS, NUM_LAYERS, BeamFullModelNvfp4Sm103Tp4 + + +def _checkpoint_config(pretrained_config) -> tuple: + """The checkpoint directory and its ``config.json`` mapping, read as stored.""" + checkpoint_dir = getattr(pretrained_config, "name_or_path", None) or getattr( + pretrained_config, "_name_or_path", None + ) + if not checkpoint_dir or not os.path.isfile(os.path.join(checkpoint_dir, "config.json")): + raise ValueError( + "the Beam NVFP4 target is built from the checkpoint's own config.json; the " + f"configuration names no readable checkpoint directory ({checkpoint_dir!r})" + ) + with open(os.path.join(checkpoint_dir, "config.json")) as f: + return checkpoint_dir, json.load(f) + + +def _log_trace(phase: str, trace: list) -> None: + """Record the stages a traced step executed, in what order, and its requests' spans.""" + layers = [r["layer"] for r in trace if r["stage"] == "layer"] + stages = [r["stage"] for r in trace] + spans = trace[0].get("requests", []) + logger.info( + f"[modeling_v2] first {phase} step: {len(spans)} requests (past, new tokens) {spans}, " + f"stages {stages[0]}..{stages[-1]}, {len(layers)} logical " + f"layers in order 0..{NUM_LAYERS - 1}: {layers == list(range(NUM_LAYERS))}, " + f"{sum(r.get('mlp') == 'sparse' for r in trace)} sparse, " + f"{sum(r.get('attention') == 'full' for r in trace)} full-attention" + ) + + +def check_drafter_config(draft_config, target_config: dict, target_layer_ids) -> None: + """Raise ValueError unless ``draft_config`` is a DFlash 2 drafter this target can serve. + + The drafter reads the target's completed block outputs at ``target_layer_ids`` (zero-based + logical layers, ascending), shares the target's vocabulary, embedding and head (no draft + to target id map), runs the DFlash 2 recipe (block convolutions and the candidate + selector) at the target's hidden size, and declares its own unary logit softcap. + """ + pretrained = draft_config.pretrained_config + dflash = getattr(pretrained, "dflash_config", None) or {} + problems = [] + if pretrained.hidden_size != target_config["hidden_size"]: + problems.append(f"hidden_size {pretrained.hidden_size} != {target_config['hidden_size']}") + vocab = target_config["vocab_size"] + for key in ("vocab_size", "draft_vocab_size"): + value = getattr(pretrained, key, vocab) + if value != vocab: + problems.append(f"{key} {value} != the target's {vocab}") + layers = list(target_layer_ids or []) + if ( + not layers + or layers != sorted(set(layers)) + or not all(type(i) is int and 0 <= i < NUM_LAYERS for i in layers) + or layers != list(dflash.get("target_layer_ids") or []) + ): + problems.append( + f"target_layer_ids {layers} must be the drafter's own ascending, distinct layers in " + f"0..{NUM_LAYERS - 1} (dflash_config: {dflash.get('target_layer_ids')})" + ) + for key in ("conv_kernel_size", "conv_group_size", "selector_rank", "selector_top_k"): + if not int(dflash.get(key, 0) or 0) > 0: + problems.append(f"dflash_config.{key} is missing") + if dflash.get("final_logit_softcapping") is None: + problems.append("dflash_config.final_logit_softcapping is missing") + if problems: + raise ValueError( + "the drafter is not a DFlash 2 drafter of this target: " + "; ".join(problems) + ) + + +class _Target: + """One engine step through all 52 layers; traces the first single- and multi-request step.""" + + phase = "" + + def __init__(self, model: BeamFullModelNvfp4Sm103Tp4): + self.model = model + self.traced = set() + + def _run(self, attn_metadata, input_ids, position_ids, logit_rows, **kwargs): + kind = "batched" if attn_metadata.num_seqs > 1 else "single-request" + trace = None if kind in self.traced else [] + out = self.model( + input_ids, position_ids, attn_metadata, logit_rows=logit_rows, trace=trace, **kwargs + ) + if trace is not None: + self.traced.add(kind) + _log_trace(f"{kind} {self.phase}", trace) + return out + + +class PrefillTarget(_Target): + """A step with a context request, possibly followed by generation requests.""" + + phase = "context" + + def forward( + self, + attn_metadata, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + *, + return_context_logits: bool, + ) -> torch.Tensor: + rows = "all" if return_context_logits else None + return self._run(attn_metadata, input_ids, position_ids, rows) + + +class DecodeTarget(_Target): + """A step of generation requests: one token each.""" + + phase = "generation" + + def forward( + self, + attn_metadata, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + *, + return_context_logits: bool, + ) -> torch.Tensor: + return self._run(attn_metadata, input_ids, position_ids, None) + + +class VerifyTarget(_Target): + """A speculative step: contexts first, then ``P`` verification rows per generation request. + + Copies the completed block output of every capture layer into the step's DFlash capture + buffer (``maybe_capture_hidden_states`` keeps only its own layers) and returns the logits + of the verification rows, which must be the engine's ``gather_ids``, with the last block's + output. + """ + + phase = "verification" + + def forward( + self, + attn_metadata, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + spec_metadata, + *, + extra_kv_tokens: int, + ) -> tuple: + rows_per_generation = spec_metadata.runtime_draft_len + 1 + if spec_metadata.runtime_tokens_per_gen_step != rows_per_generation: + raise ValueError( + f"the step verifies {spec_metadata.runtime_tokens_per_gen_step} rows per " + f"generation request but drafts {spec_metadata.runtime_draft_len}; a linear " + "draft verifies its proposals and the pending token" + ) + return self._run( + attn_metadata, + input_ids, + position_ids, + "verification", + generation_rows=rows_per_generation, + extra_kv_tokens=extra_kv_tokens, + block_output=spec_metadata.maybe_capture_hidden_states, + gather_ids=spec_metadata.gather_ids, + return_hidden=True, + ) + + +class ModelingV2BeamFullModelNvfp4Sm103Tp4(nn.Module): + """The engine-facing Beam NVFP4 full model (one rank of four).""" + + def __init__(self, model_config: ModelConfig): + super().__init__() + self.model_config = model_config + mapping = model_config.mapping + self.checkpoint_dir, checkpoint_config = _checkpoint_config(model_config.pretrained_config) + self.model = BeamFullModelNvfp4Sm103Tp4( + checkpoint_config, + mapping, + torch.device("cuda", torch.cuda.current_device()), + max_tokens=MAX_NUM_TOKENS, + ) + self._prefill = PrefillTarget(self.model) + self._decode = DecodeTarget(self.model) + self.load_record: Optional[dict] = None + # The native one-model pieces, built only for a speculative config. + self.spec_config = getattr(model_config, "spec_config", None) + self.draft_config: Optional[ModelConfig] = None + self.draft_model: Optional[nn.Module] = None + self.spec_worker: Optional[nn.Module] = None + self.use_separate_draft_kv_cache = False + self.extra_kv_tokens = 0 + self.draft_load_record: Optional[dict] = None + if self.spec_config is not None: + self._build_dflash(model_config, checkpoint_config) + logger.info( + f"[modeling_v2] {type(self).__name__}: {NUM_LAYERS} layers, rank {mapping.rank} of " + f"{mapping.world_size}, max {MAX_NUM_TOKENS} tokens per call, speculative decoding " + f"{'DFlash' if self.spec_worker is not None else 'off'}" + ) + + def _build_dflash(self, model_config: ModelConfig, checkpoint_config: dict) -> None: + """Compose the native DFlash drafter, worker and draft-cache choice on this target.""" + from tensorrt_llm._torch.models.modeling_dflash import ( + DFlashSoftcapLogitsProcessor, + dflash_final_logit_softcap, + ) + from tensorrt_llm._torch.models.modeling_speculative import ( + _set_draft_kv_cache_quant_algo, + external_drafter_config_kwargs, + get_draft_model, + ) + from tensorrt_llm._torch.speculative import ( + get_num_extra_kv_tokens, + get_spec_worker, + should_use_separate_draft_kv_cache, + ) + + spec_config = self.spec_config + if not spec_config.spec_dec_mode.is_dflash(): + raise ValueError( + f"the Beam NVFP4 target composes a DFlash drafter only, got " + f"{type(spec_config).__name__}" + ) + self._verify = VerifyTarget(self.model) + self.draft_config = ModelConfig.from_pretrained( + spec_config.speculative_model, + **external_drafter_config_kwargs(model_config, spec_config), + ) + _set_draft_kv_cache_quant_algo(self.draft_config, model_config) + self.draft_config.extra_attrs = model_config.extra_attrs + check_drafter_config(self.draft_config, checkpoint_config, spec_config.target_layer_ids) + # The draft pages live in the engine's separate draft KV cache manager, which the + # worker reads through its own pool views (TRTLLM draft attention). + self.use_separate_draft_kv_cache = should_use_separate_draft_kv_cache(spec_config) + self.draft_model = get_draft_model(model_config, self.draft_config, None, self) + self.draft_model.logits_processor = DFlashSoftcapLogitsProcessor( + dflash_final_logit_softcap(self.draft_config.pretrained_config) + ) + self.spec_worker = get_spec_worker( + spec_config, + model_config, + model_config.mapping, + use_separate_draft_kv_cache=self.use_separate_draft_kv_cache, + ) + if self.spec_worker is None: + raise ValueError(f"no speculative worker serves {spec_config.spec_dec_mode!r}") + self.spec_worker.set_draft_model(self.draft_model) + self.extra_kv_tokens = get_num_extra_kv_tokens(spec_config) + + @property + def config(self): + return self.model_config.pretrained_config + + @property + def vocab_size_padded(self) -> int: + return self.model.vocab_size + + def load_checkpoint(self, checkpoint_dir: str) -> dict: + """Fill every parameter from ``checkpoint_dir`` (collective; see the component's load).""" + if os.path.realpath(checkpoint_dir) != os.path.realpath(self.checkpoint_dir): + raise ValueError( + f"weights from {checkpoint_dir!r} do not belong to the configuration of " + f"{self.checkpoint_dir!r}" + ) + self.load_record = self.model.load(checkpoint_dir) + return self.load_record + + def shared_vocabulary_modules(self) -> tuple: + """Native ``Embedding`` and vocabulary-parallel ``LMHead`` over the target's own slices. + + Each module's weight is the component's registered rank-local parameter itself + (``model.embed_tokens.weight`` / ``lm_head.weight``, rows ``[rank * V/4, (rank + 1) * + V/4)``), not a copy: the embedding masks ids outside the slice and all-reduces, the + head projects to the rank's vocabulary shard without gathering or capping. Requires a + completed target load. + """ + from tensorrt_llm._torch.modules.embedding import Embedding, LMHead + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + if not self.model.usable: + raise RuntimeError("the target's embedding and head are shared after its load") + mapping = self.model_config.mapping + vocab, hidden = self.model.vocab_size, self.model.hidden + modules = ( + Embedding( + vocab, + hidden, + dtype=torch.bfloat16, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=False, + ), + LMHead( + vocab, + hidden, + dtype=torch.bfloat16, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=False, + ), + ) + lo, hi = self.model.vocab_interval + for module, name in zip(modules, (EMBED, HEAD)): + weight = self.model.tensor(name) + if tuple(module.weight.shape) != tuple(weight.shape): + raise ValueError( + f"{type(module).__name__} holds {tuple(module.weight.shape)} rows per rank; " + f"the target's {name} slice is {tuple(weight.shape)}" + ) + module.weight = weight + embed, head = modules + if (embed.vocab_start_index, embed.vocab_end_index) != (lo, hi): + raise ValueError( + f"the embedding owns ids [{embed.vocab_start_index}, {embed.vocab_end_index}) on " + f"this rank; the target's slice is [{lo}, {hi})" + ) + return embed, head + + def load_draft_weights(self, weights: dict, weight_mapper=None) -> None: + """Load the drafter's own tensors, then give it the target's embedding and head. + + The drafter checkpoint carries neither (``DFlashForCausalLM.WEIGHTS_SHARED_WITH_TARGET``) + and its own loader rejects every other missing tensor; the shared modules are views of + the target's slices (``shared_vocabulary_modules``), handed over through the drafter's + native ``load_weights_from_target_model``. + """ + if self.draft_model is None: + raise RuntimeError("this target was built without a drafter") + self.draft_model.load_weights(weights=weights, weight_mapper=weight_mapper) + embed, head = self.shared_vocabulary_modules() + self.draft_model.load_weights_from_target_model( + types.SimpleNamespace(model=types.SimpleNamespace(embed_tokens=embed), lm_head=head) + ) + draft_embed = self.draft_model.draft_model_full.model.embed_tokens + self.draft_load_record = { + "tensors": len(weights), + "embed_shares_target_storage": draft_embed.weight.data_ptr() + == self.model.tensor(EMBED).data_ptr(), + "head_shares_target_storage": self.draft_model.lm_head.weight.data_ptr() + == self.model.tensor(HEAD).data_ptr(), + "vocab_interval": list(self.model.vocab_interval), + } + logger.info(f"[modeling_v2] DFlash drafter loaded: {self.draft_load_record}") + + def forward( + self, + attn_metadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + spec_metadata=None, + resource_manager=None, + **kwargs, + ): + if inputs_embeds is not None or input_ids is None: + raise ValueError("the Beam NVFP4 target takes token ids, not input embeddings") + if kwargs.get("lora_params") is not None: + raise ValueError("LoRA is not supported by this target") + ids = input_ids.reshape(-1).to(torch.int64) + if self.spec_worker is None: + if spec_metadata is not None: + raise ValueError( + "speculative metadata reached a Beam target built without a speculative config" + ) + target = self._decode if attn_metadata.num_contexts == 0 else self._prefill + return target.forward( + attn_metadata, ids, position_ids, return_context_logits=return_context_logits + ) + if spec_metadata is None: + raise ValueError( + "the DFlash target verifies every step; the engine passed no speculative metadata" + ) + logits, hidden = self._verify.forward( + attn_metadata, ids, position_ids, spec_metadata, extra_kv_tokens=self.extra_kv_tokens + ) + return self.spec_worker( + input_ids=ids, + position_ids=position_ids, + hidden_states=hidden, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) + + +__all__ = [ + "DecodeTarget", + "ModelingV2BeamFullModelNvfp4Sm103Tp4", + "PrefillTarget", + "VerifyTarget", + "check_drafter_config", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/serving.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/serving.py new file mode 100644 index 000000000000..e5f9247cda6e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/serving.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Text serving for the Beam NVFP4 full-model target: the channel parser and a native launch. + +The model writes its reply as channel messages after the chat template's bare +``<|start|>assistant`` generation prefix:: + + <|channel|>analysis<|message|>...<|end|><|start|>assistant<|channel|>final<|message|>...<|return|> + +The delimiters are ordinary (non-special) vocabulary tokens, so they reach the detokenized text +whatever ``skip_special_tokens`` says. ``BeamHarmonyTextParser`` (reasoning parser key +``beam_harmony_text``) routes analysis bodies to ``reasoning_content`` and final bodies to +``content`` and drops every header and terminator. A reply without channel headers is final +text. It is incremental: a possible delimiter split across stream chunks is held back until it +is resolved, and ``finish`` releases a held tail that is ordinary text. ``parse`` runs the same +machine over a whole reply, so streamed and non-streamed fields agree. Tool calls are not +supported; a message on any other channel is ignored. + +``main`` launches the existing OpenAI-compatible server (``launch_server``) on the supported +deployment: four ranks (attention TP4, MoE TP4/EP1, PP1), ``modeling_v2="require"``, TRTLLM +attention, ``KVCacheManagerV2`` with block reuse, the native TorchSampler, eager execution +(no CUDA graphs, no overlap scheduler, no chunked prefill) and one frontend. The native +scheduler batches up to ``--max-batch-size`` requests per engine step (default 1: one sequence +at a time). It owns no scheduler, sampler, route or evaluator:: + + python -m tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.serving \\ + --checkpoint [--draft-checkpoint ] [--port 8000] [--max-seq-len 69632] \\ + [--max-num-tokens 8192] [--max-batch-size 1] [--iter-stats] + +``--draft-checkpoint`` adds DFlash speculative decoding with that DFlash 2 drafter: the native +one-model DFlash worker with ``max_draft_len`` 7 (a verification block of eight rows), TRTLLM +draft attention over the engine's paged draft KV pool and strict greedy verification; the +one-model speculative sampler replaces the TorchSampler and only target-verified tokens reach +the client. ``--iter-stats`` enables the engine's iteration statistics (``GET /metrics``: +per-step scheduled, context and generation request counts, latency and memory). +""" + +from __future__ import annotations + +import argparse +import os +from typing import Any, Optional + +from tensorrt_llm.llmapi.reasoning_parser import ( + BaseReasoningParser, + ReasoningParserResult, + register_reasoning_parser, +) + +START, END, MESSAGE, CHANNEL, CONSTRAIN, RETURN, CALL = ( + "<|start|>", + "<|end|>", + "<|message|>", + "<|channel|>", + "<|constrain|>", + "<|return|>", + "<|call|>", +) +TAGS = (START, END, MESSAGE, CHANNEL, CONSTRAIN, RETURN, CALL) +ANALYSIS, FINAL = "analysis", "final" +PARSER_KEY = "beam_harmony_text" + + +def _held_tail(text: str) -> int: + """Length of the longest suffix of ``text`` that is a proper prefix of a delimiter.""" + for n in range(min(len(text), max(len(t) for t in TAGS) - 1), 0, -1): + tail = text[-n:] + if any(t.startswith(tail) and t != tail for t in TAGS): + return n + return 0 + + +@register_reasoning_parser(PARSER_KEY) +class BeamHarmonyTextParser(BaseReasoningParser): + """Split a Beam channel reply into ``reasoning_content`` (analysis) and ``content`` (final).""" + + needs_raw_special_tokens = True + reasoning_start = CHANNEL + ANALYSIS + MESSAGE + reasoning_end = END + + def __init__(self, *, chat_template_kwargs: Optional[dict[str, Any]] = None) -> None: + super().__init__(chat_template_kwargs=chat_template_kwargs) + self._reset() + + def _reset(self) -> None: + # "header": before a channel's <|message|>; "body": inside a message; "between": after + # <|end|> until the next header; "done": after <|return|> or <|call|>. + self._state = "header" + self._channel: Optional[str] = None + self._header = "" + self._buffer = "" + self._saw_header = False + + def _emit(self, text: str, out: list) -> None: + if not text: + return + if self._channel == ANALYSIS: + out[1].append(text) + elif self._channel == FINAL or (self._channel is None and not self._saw_header): + out[0].append(text) + + def _feed(self, text: str) -> ReasoningParserResult: + out = ([], []) + data = self._buffer + text + self._buffer = "" + while data: + if self._state == "done": + data = "" + break + hits = [(data.find(t), t) for t in TAGS if t in data] + pos, tag = min(hits) if hits else (-1, None) + if tag is None: + hold = _held_tail(data) + chunk, self._buffer = data[: len(data) - hold], data[len(data) - hold :] + data = "" + else: + chunk, data = data[:pos], data[pos + len(tag) :] + if self._state == "body": + self._emit(chunk, out) + elif self._state == "header": + if self._header or self._saw_header or tag in (CHANNEL, MESSAGE, CONSTRAIN): + self._header += chunk + else: + # Text before any header: a reply without channels is final text. + self._emit(chunk, out) + if tag is None: + break + if tag == CHANNEL: + self._state, self._header, self._saw_header = "header", "", True + elif tag == MESSAGE: + if self._state == "header": + self._channel = self._header.split(CONSTRAIN)[0].strip().split(" ")[0] or None + self._header = "" + self._state = "body" + elif tag == CONSTRAIN: + self._header += CONSTRAIN + elif tag in (END, START): + self._state, self._channel, self._header = "between", None, "" + elif tag in (RETURN, CALL): + self._state, self._channel = "done", None + return ReasoningParserResult(content="".join(out[0]), reasoning_content="".join(out[1])) + + def parse(self, text: str) -> ReasoningParserResult: + self._reset() + result = self._feed(text) + tail = self.finish() + return ReasoningParserResult( + content=result.content + tail.content, + reasoning_content=result.reasoning_content + tail.reasoning_content, + ) + + def parse_delta(self, delta_text: str) -> ReasoningParserResult: + return self._feed(delta_text) + + def finish(self) -> ReasoningParserResult: + held, self._buffer = self._buffer, "" + out = ([], []) + if self._state == "body" and held and held not in TAGS: + self._emit(held, out) + elif self._state == "header" and not self._saw_header and not self._header and held: + self._emit(held, out) + return ReasoningParserResult(content="".join(out[0]), reasoning_content="".join(out[1])) + + +#: The supported deployment's settings (one source for the launcher and its records). +DEPLOYMENT = { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 1, + "moe_tensor_parallel_size": 4, + "moe_expert_parallel_size": 1, + "enable_attention_dp": False, + "attn_backend": "TRTLLM", + "modeling_v2": "require", + "disable_overlap_scheduler": True, + "enable_chunked_prefill": False, +} + + +def load_tokenizer(checkpoint: str): + """The checkpoint's fast tokenizer, loaded by its concrete class. + + ``AutoTokenizer`` would first parse config.json as a generic transformers configuration, + which rejects Beam's per-layer-type RoPE parameters; the tokenizer itself needs none of it. + """ + from transformers import PreTrainedTokenizerFast + + return PreTrainedTokenizerFast.from_pretrained(checkpoint) + + +def dflash_config(draft_checkpoint: str): + """The DFlash speculative config of the supported deployment for ``draft_checkpoint``. + + ``max_draft_len`` is the one verification width the target serves (``routing``); the + capture layers and mask token come from the drafter's own ``dflash_config``. + """ + from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig + + from ..routing import DFLASH_MAX_DRAFT_LEN + + if not os.path.isfile(os.path.join(draft_checkpoint, "config.json")): + raise ValueError(f"the drafter checkpoint {draft_checkpoint!r} has no config.json") + return DFlashDecodingConfig( + max_draft_len=DFLASH_MAX_DRAFT_LEN, + speculative_model=draft_checkpoint, + attention_backend="TRTLLM", + ) + + +def llm_args( + checkpoint: str, + *, + max_seq_len: int = 69632, + max_num_tokens: int = 8192, + max_batch_size: int = 1, + kv_cache_max_tokens: Optional[int] = None, + iter_stats: bool = False, + draft_checkpoint: Optional[str] = None, +) -> dict: + """Native LLM API arguments of the supported deployment for ``checkpoint``. + + ``max_batch_size`` bounds the requests the native scheduler runs in one engine step; + ``kv_cache_max_tokens`` (default ``2 * max_seq_len``) sizes the shared KV cache. + ``draft_checkpoint`` adds DFlash speculative decoding with that drafter + (``dflash_config``); its draft KV pool is the engine's separate draft manager. + """ + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + + from .model import MAX_NUM_TOKENS + from .weights import BeamNvfp4CheckpointLoader + + if not 1 <= max_num_tokens <= MAX_NUM_TOKENS: + raise ValueError(f"max_num_tokens must be in [1, {MAX_NUM_TOKENS}], got {max_num_tokens}") + args = { + "model": checkpoint, + "tokenizer": load_tokenizer(checkpoint), + "backend": "pytorch", + **DEPLOYMENT, + "max_batch_size": max_batch_size, + "max_seq_len": max_seq_len, + "max_num_tokens": max_num_tokens, + "cuda_graph_config": None, + "kv_cache_config": KvCacheConfig( + enable_block_reuse=True, + enable_partial_reuse=False, + use_kv_cache_manager_v2=True, + max_tokens=kv_cache_max_tokens or 2 * max_seq_len, + ), + "checkpoint_loader": BeamNvfp4CheckpointLoader(), + "reasoning_parser": PARSER_KEY, + "enable_iter_perf_stats": iter_stats, + } + if draft_checkpoint is not None: + args["speculative_config"] = dflash_config(draft_checkpoint) + return args + + +def main(argv: Optional[list] = None) -> None: + ap = argparse.ArgumentParser(description="Serve the Beam NVFP4 full model (OpenAI API).") + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--host", default="127.0.0.1") + ap.add_argument("--port", type=int, default=8000) + ap.add_argument("--served-model-name", default=None) + ap.add_argument( + "--chat-template", default=None, help="default: /chat_template.jinja" + ) + ap.add_argument("--max-seq-len", type=int, default=69632) + ap.add_argument("--max-num-tokens", type=int, default=8192) + ap.add_argument("--max-batch-size", type=int, default=1) + ap.add_argument("--kv-cache-max-tokens", type=int, default=None) + ap.add_argument("--iter-stats", action="store_true") + ap.add_argument( + "--draft-checkpoint", + default=None, + help="DFlash 2 drafter directory; without it, speculation is off", + ) + args = ap.parse_args(argv) + + from tensorrt_llm.commands.serve import launch_server + + # The capacity is the explicit max_tokens above; no estimation pass on a dry-run cache. + os.environ.setdefault("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") + template = args.chat_template or os.path.join(args.checkpoint, "chat_template.jinja") + launch_server( + args.host, + args.port, + llm_args( + args.checkpoint, + max_seq_len=args.max_seq_len, + max_num_tokens=args.max_num_tokens, + max_batch_size=args.max_batch_size, + kv_cache_max_tokens=args.kv_cache_max_tokens, + iter_stats=args.iter_stats, + draft_checkpoint=args.draft_checkpoint, + ), + chat_template=template, + served_model_name=args.served_model_name, + multi_frontend_enabled=False, + ) + + +if __name__ == "__main__": + main() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/weights.py new file mode 100644 index 000000000000..9199b09237f4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/weights.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weights of the Beam NVFP4 full-model target: a direct preload through the native loader. + +``BeamNvfp4CheckpointLoader`` is passed as the LLM API's per-instance ``checkpoint_loader``. +The native ``ModelLoader`` constructs the target, materializes it on the rank's device and then +calls ``load_weights(checkpoint_dir, mapping, model=...)``; this loader hands the directory to +the target's ``load_checkpoint``, which streams bounded tiles straight into the target's own +rank-local storage (packed E2M1 codes, E4M3 block scales, FP32 scalars, BF16/FP32 dense +tensors) as one all-rank transaction, and returns no weight dictionary. The native lifecycle +continues as for any preloaded checkpoint (``is_weights_preloaded``). Configuration loading is +the HF loader's; everything else is the HF checkpoint loader's own behavior. + +With a DFlash drafter the native loader then reads the drafter checkpoint +(``speculative_model``) through the same loader, without a model: that second load is the HF +loader's ordinary weight dictionary, which the target's ``load_draft_weights`` hands to the +drafter's own loader. It is accepted only after the target transaction completed and only for +a directory other than the target's, so neither load can run the other's path; the target +stays preloaded. +""" + +from __future__ import annotations + +import os +from typing import Any, Optional + +from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import HfCheckpointLoader +from tensorrt_llm.mapping import Mapping + + +class BeamNvfp4CheckpointLoader(HfCheckpointLoader): + """Preload the Beam NVFP4 full-model target; read a drafter's weights after it.""" + + def __init__(self): + super().__init__() + self._preloaded = False + self._target_dir: Optional[str] = None + + def __reduce__(self): + # Workers rebuild the loader from its class: it carries no state before a load. + return (type(self), ()) + + def load_weights(self, checkpoint_dir: str, mapping: Mapping, **kwargs) -> dict[str, Any]: + from .modeling import ModelingV2BeamFullModelNvfp4Sm103Tp4 + + model = kwargs.get("model") + if model is None: + return self._load_draft_weights(checkpoint_dir, mapping, **kwargs) + if not isinstance(model, ModelingV2BeamFullModelNvfp4Sm103Tp4): + raise TypeError( + "BeamNvfp4CheckpointLoader loads only ModelingV2BeamFullModelNvfp4Sm103Tp4 " + f"(modeling_v2='require'); got {type(model).__name__}" + ) + self._preloaded = False + self._target_dir = None + model.load_checkpoint(checkpoint_dir) + self._target_dir = os.path.realpath(checkpoint_dir) + self._preloaded = True + return {} + + def _load_draft_weights(self, checkpoint_dir: str, mapping: Mapping, **kwargs) -> dict: + """The drafter checkpoint's tensors, read by the HF loader for the drafter's own load.""" + if not self._preloaded: + raise RuntimeError( + f"{checkpoint_dir!r}: a weight load without the target model is the drafter's, " + "which follows the Beam target's own load" + ) + if os.path.realpath(checkpoint_dir) == self._target_dir: + raise ValueError( + f"{checkpoint_dir!r} is the Beam target's checkpoint; a drafter load names the " + "drafter checkpoint (speculative_model)" + ) + return super().load_weights(checkpoint_dir, mapping=mapping, **kwargs) + + def is_weights_preloaded(self) -> bool: + return self._preloaded + + +__all__ = ["BeamNvfp4CheckpointLoader"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/__init__.py new file mode 100644 index 000000000000..e6278158bae4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE feed-forward of decoder block 3 / sm_103 / tp1, as a directly constructed component.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/moe.py new file mode 100644 index 000000000000..626143c4acda --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/moe.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE feed-forward of the sparse decoder blocks 3-5 / sm_103 / tp1, constructed directly. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine builds this class, and it serves no request. +A caller constructs it from the checkpoint's configuration, loads the MoE's 770 tensors +and passes it the already-normalized hidden states that enter its sparse block's +feed-forward sublayer; ``layer_idx`` (3, 4 or 5, as the router accepts) names that block, +and each block owns its own router and experts. It composes three accepted components and +adds no arithmetic of its own:: + + weights, ids = router(hidden) ../moe_router_bf16__sm_103__tp1 + out = dispatch_combine(hidden, ids, weights, ../moe_dispatch_combine_bf16__sm_103__tp1 + lambda e, rows: experts[e](rows)) + ../moe_expert_bf16__sm_103__tp1, 256 of them + +That is, as the source model computes it: fp32 grouped sigmoid top-k routing with a +selection-only correction bias (8 of 256 experts per token, from 4 of 8 groups) and +unbiased mixing weights; each selected expert, in ascending id order, runs its own bf16 +SwiGLU on its tokens' rows; each expert output times its fp32 weight is rounded to bf16, +the routes are summed in fp32, scaled by the routed scaling factor 2.5 once and rounded +to bf16. No residual, normalization or collective is applied here. + +Host synchronization: routing runs on the GPU without one. Dispatch reads the selected +experts and each expert's tokens back to the host -- two synchronizations for the expert +list and two per selected expert -- so a call is not capturable in a CUDA graph. Token +rows, products, roundings, sums and the scaling run on the GPU. +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import Mapping + +import torch +from torch import nn + +from ..moe_dispatch_combine_bf16__sm_103__tp1 import dispatch_combine +from ..moe_expert_bf16__sm_103__tp1.expert import WEIGHT_SHAPES as EXPERT_WEIGHT_SHAPES +from ..moe_expert_bf16__sm_103__tp1.expert import BeamMoeExpertBf16Sm103Tp1 +from ..moe_router_bf16__sm_103__tp1.router import BeamMoeRouterBf16Sm103Tp1 +from ..moe_router_bf16__sm_103__tp1.router import check_config as check_router_config + +#: Fields of the checkpoint's configuration this component reads beyond the router's. +SUPPORTED_CONFIG = { + "moe_intermediate_size": 2560, + "hidden_act": "silu", + "routed_scaling_factor": 2.5, +} +SUPPORTED_LAYER_IDX = 3 + +_SM = (10, 3) + + +def check_config(config, layer_idx: int) -> None: + """Raise ValueError naming the first field outside what this component computes.""" + check_router_config(config, layer_idx) + for field, expected in SUPPORTED_CONFIG.items(): + actual = getattr(config, field, None) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this MoE computes {expected!r}") + + +def weight_shapes(config) -> "OrderedDict[str, tuple[tuple[int, ...], torch.dtype]]": + """The MoE's tensors by checkpoint name (relative to ``mlp.``), in order, with dtypes.""" + shapes = OrderedDict( + [ + ("gate.weight", ((config.num_experts, config.hidden_size), torch.float32)), + ("gate.e_score_correction_bias", ((config.num_experts,), torch.float32)), + ] + ) + for expert_id in range(config.num_experts): + for name, shape in EXPERT_WEIGHT_SHAPES.items(): + shapes[f"experts.{expert_id}.{name}"] = (shape, torch.bfloat16) + return shapes + + +def check_weights(expected: Mapping, weights: Mapping[str, torch.Tensor], what: str) -> None: + """Raise unless ``weights`` holds exactly ``expected``'s names, dtypes and shapes. + + Reads tensor metadata only, never values, and copies nothing. + """ + if not isinstance(weights, Mapping): + raise TypeError(f"weights must be a mapping, got {type(weights).__name__}") + missing = [n for n in expected if n not in weights] + unexpected = sorted(str(n) for n in weights if n not in expected) + if missing or unexpected: + raise ValueError( + f"{what} weights mismatch: missing {missing[:8]}{'...' if len(missing) > 8 else ''}, " + f"unexpected {unexpected[:8]}{'...' if len(unexpected) > 8 else ''}" + ) + for name, (shape, dtype) in expected.items(): + tensor = weights[name] + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + if tensor.dtype != dtype: + raise TypeError(f"{name} must be {dtype}, got {tensor.dtype}") + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must be {list(shape)}, got {list(tensor.shape)}") + if tensor.layout != torch.strided or tensor.device.type not in ("cpu", "cuda"): + raise ValueError( + f"{name} must be a dense host or CUDA tensor, got {tensor.layout} on {tensor.device}" + ) + + +class BeamMoeBf16Sm103Tp1(nn.Module): + """A sparse block's MoE: normalized BF16 hidden states in, the combined BF16 MoE output out. + + State: ``gate``, the router (FP32 ``weight [256, 5120]`` and selection-only + ``e_score_correction_bias [256]``), and ``experts``, 256 routed experts, each with its + own bias-free BF16 ``gate_proj``/``up_proj [2560, 5120]`` and ``down_proj [5120, 2560]``. + All are loaded together with ``load_weights``. + + ``forward(hidden_states)`` takes ``[T, 5120]`` ``torch.bfloat16`` on the MoE's device, + T >= 1 tokens already normalized by the caller, and returns ``[T, 5120]`` + ``torch.bfloat16`` on that device, row ``t`` for token ``t``. Inputs outside that are + rejected before any kernel runs, and the input is never modified. Measured at T = 1, 4, + 17 and 129. + """ + + def __init__( + self, config, layer_idx: int = SUPPORTED_LAYER_IDX, device: str | torch.device = "cuda" + ): + super().__init__() + check_config(config, layer_idx) + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this MoE is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self._shapes = weight_shapes(config) + self.gate = BeamMoeRouterBf16Sm103Tp1(config, layer_idx, device=device) + self.experts = nn.ModuleList( + BeamMoeExpertBf16Sm103Tp1(device=device) for _ in range(config.num_experts) + ) + self._loaded = False + + def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": + """The MoE's tensors by checkpoint name (relative to ``mlp.``), in checkpoint order.""" + tensors = OrderedDict( + [ + ("gate.weight", self.gate.weight), + ("gate.e_score_correction_bias", self.gate.e_score_correction_bias), + ] + ) + for expert_id, expert in enumerate(self.experts): + for name, tensor in expert.checkpoint_tensors().items(): + tensors[f"experts.{expert_id}.{name}"] = tensor + return tensors + + def check_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Raise unless ``weights`` is exactly the MoE's 770 tensors by name, dtype and shape.""" + check_weights(self._shapes, weights, "MoE") + + @torch.no_grad() + def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Copy exactly the MoE's 770 tensors, by checkpoint name relative to ``mlp.``. + + ``gate.weight`` and ``gate.e_score_correction_bias`` are FP32; each + ``experts..{gate,up,down}_proj.weight`` is BF16 and loads into expert ``e``. + Values are copied as given, [out, in], without transposing or converting; host + tensors are copied to the MoE's device. Every tensor is checked before any + parameter changes, so a rejected load leaves the MoE as it was. The MoE runs only + after a load completed: a failure while copying leaves it unloaded. + """ + self.check_weights(weights) + self._loaded = False + self.gate.load_weights(weights["gate.weight"], weights["gate.e_score_correction_bias"]) + for expert_id, expert in enumerate(self.experts): + prefix = f"experts.{expert_id}." + expert.load_weights({name: weights[prefix + name] for name in EXPERT_WEIGHT_SHAPES}) + self._loaded = True + + def _check_input(self, hidden_states: torch.Tensor) -> None: + """Raise on anything this MoE would otherwise compute wrongly.""" + if not self._loaded: + raise RuntimeError("load_weights must be called before the MoE runs") + if not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"hidden_states must be a torch.Tensor, got {type(hidden_states).__name__}" + ) + if hidden_states.dtype != torch.bfloat16: + raise TypeError(f"hidden_states must be torch.bfloat16, got {hidden_states.dtype}") + if hidden_states.layout != torch.strided: + raise ValueError(f"hidden_states must be a dense tensor, got {hidden_states.layout}") + if hidden_states.dim() != 2 or hidden_states.shape[1] != self.hidden_size: + raise ValueError( + f"hidden_states must be [tokens, {self.hidden_size}], " + f"got {list(hidden_states.shape)}" + ) + if hidden_states.shape[0] < 1: + raise ValueError("hidden_states must hold at least one token") + if hidden_states.device != self.gate.weight.device: + raise ValueError( + f"hidden_states must be on the MoE's device {self.gate.weight.device}, " + f"got {hidden_states.device}" + ) + + def _expert(self, expert_id: int, rows: torch.Tensor) -> torch.Tensor: + """The dispatch provider: expert ``expert_id``'s own unweighted output for ``rows``.""" + return self.experts[expert_id](rows) + + @torch.no_grad() + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self._check_input(hidden_states) + mixing_weights, expert_ids = self.gate(hidden_states) + return dispatch_combine.beam_dispatch_combine( + hidden_states, expert_ids, mixing_weights, self._expert + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/__init__.py new file mode 100644 index 000000000000..a0ca5d84919b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE dispatch and weighted combination of decoder block 3 (H5120 / E256 / top-8) / sm_103 / tp1.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/dispatch_combine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/dispatch_combine.py new file mode 100644 index 000000000000..62bc3cc9f6ad --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/dispatch_combine.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE dispatch and weighted combination of decoder block 3 / sm_103 / tp1, called directly. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine calls this function, and it serves no request. +A caller passes the hidden states that enter block 3's feed-forward sublayer, the router's +result for them -- eight expert ids and eight mixing weights per token -- and an expert +provider. The function sends each token's rows to its experts through the provider and +returns the weighted, combined and scaled expert outputs. It holds no state: no router +parameter, no expert matrix and no configuration object, so there is nothing to build or +load before calling it. It does not call the built-in Beam modeling. + +Ownership. This component owns dispatch and combination; the provider owns the expert +computation. ``expert_provider(expert_id, rows)`` is called once for every selected +expert, in ascending id order, with that expert's ``rows`` -- ``[n, 5120]`` +``torch.bfloat16``, ``n >= 1``, the routed tokens in ascending token order -- and returns +the expert's unweighted output for them: ``[n, 5120]`` ``torch.bfloat16`` on the same +device, row ``i`` for input row ``i``. It applies no mixing weight and no routed scaling +factor; both are applied here. A routed expert component per id adapts to it as +``lambda expert_id, rows: experts[expert_id](rows)``. The router owns selection: no expert +is re-selected, no bias is added and no weight is re-normalized here. No residual, +normalization or collective is applied. + +The combination, as the source model computes it (``ops/moe/beam_dispatch_combine.py``):: + + for each selected expert e, ascending: + out = expert_provider(e, hidden[tokens routed to e]) [n, 5120] bf16 + route = (out.float() * weights[:, None]).to(bfloat16) [n, 5120] each route rounded once + acc[tokens] += route.float() [T, 5120] fp32 + output = (acc * 2.5).to(bfloat16) [T, 5120] bf16 +""" + +from __future__ import annotations + +import functools +from typing import Callable + +import torch + +from .ops.moe.beam_dispatch_combine import beam_dispatch_combine_op + +HIDDEN_SIZE = 5120 +NUM_EXPERTS = 256 +EXPERTS_PER_TOKEN = 8 +ROUTED_SCALING_FACTOR = 2.5 + +#: ``expert_provider(expert_id, rows) -> output``; see the module docstring. +ExpertProvider = Callable[[int, torch.Tensor], torch.Tensor] + +_SM = (10, 3) + + +def _check_tensor(name: str, tensor, dtype: torch.dtype) -> None: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + if tensor.dtype != dtype: + raise TypeError(f"{name} must be {dtype}, got {tensor.dtype}") + if tensor.layout != torch.strided: + raise ValueError(f"{name} must be a dense tensor, got {tensor.layout}") + + +def _check_inputs( + hidden_states: torch.Tensor, + expert_ids: torch.Tensor, + mixing_weights: torch.Tensor, + expert_provider: ExpertProvider, +) -> None: + """Raise on any input this component would otherwise compute wrongly. + + Reads tensor metadata only. The ids' values are not checked: each in ``[0, 256)`` and + the eight ids of a row distinct are the caller's preconditions -- the router's output + satisfies both -- because checking them would copy the routing table to the host on + every call. + """ + _check_tensor("hidden_states", hidden_states, torch.bfloat16) + _check_tensor("expert_ids", expert_ids, torch.int64) + _check_tensor("mixing_weights", mixing_weights, torch.float32) + if hidden_states.ndim != 2 or hidden_states.shape[1] != HIDDEN_SIZE: + raise ValueError( + f"hidden_states must be [tokens, {HIDDEN_SIZE}], got {list(hidden_states.shape)}" + ) + tokens = hidden_states.shape[0] + if tokens < 1: + raise ValueError("hidden_states must hold at least one token") + for name, tensor in (("expert_ids", expert_ids), ("mixing_weights", mixing_weights)): + if tuple(tensor.shape) != (tokens, EXPERTS_PER_TOKEN): + raise ValueError( + f"{name} must be [tokens, {EXPERTS_PER_TOKEN}] = [{tokens}, " + f"{EXPERTS_PER_TOKEN}] for {tokens} tokens, got {list(tensor.shape)}" + ) + device = hidden_states.device + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"hidden_states must be on an sm_{_SM[0]}{_SM[1]} CUDA device, got {device}" + ) + for name, tensor in (("expert_ids", expert_ids), ("mixing_weights", mixing_weights)): + if tensor.device != device: + raise ValueError( + f"{name} must be on hidden_states' device {device}, got {tensor.device}" + ) + if not callable(expert_provider): + raise TypeError(f"expert_provider must be callable, got {type(expert_provider).__name__}") + + +def _checked_expert_output( + expert_provider: ExpertProvider, expert_id: int, rows: torch.Tensor +) -> torch.Tensor: + """Call the provider once and return its output unchanged, or raise if it breaks the contract.""" + out = expert_provider(expert_id, rows) + if not isinstance(out, torch.Tensor): + raise TypeError( + f"expert_provider({expert_id}) must return a torch.Tensor, got {type(out).__name__}" + ) + if out.dtype != torch.bfloat16: + raise TypeError(f"expert_provider({expert_id}) must return torch.bfloat16, got {out.dtype}") + if out.shape != rows.shape: + raise ValueError( + f"expert_provider({expert_id}) must return {list(rows.shape)} for " + f"{rows.shape[0]} rows, got {list(out.shape)}" + ) + if out.device != rows.device or out.layout != torch.strided: + raise ValueError( + f"expert_provider({expert_id}) must return a dense tensor on {rows.device}, " + f"got {out.layout} on {out.device}" + ) + return out + + +@torch.no_grad() +def beam_dispatch_combine( + hidden_states: torch.Tensor, + expert_ids: torch.Tensor, + mixing_weights: torch.Tensor, + expert_provider: ExpertProvider, +) -> torch.Tensor: + """Dispatch tokens to their experts and return the combined output. + + ``hidden_states`` ``[T, 5120]`` ``torch.bfloat16`` on an sm_103 CUDA device, T >= 1; + ``expert_ids`` ``[T, 8]`` ``torch.int64`` and ``mixing_weights`` ``[T, 8]`` + ``torch.float32`` on the same device, as the router returns them; ``expert_provider`` + as in the module docstring. Returns ``[T, 5120]`` ``torch.bfloat16`` on that device, + row ``t`` for token ``t``. Inputs outside that, and a provider output outside its + contract, are rejected with an error naming the argument; the three tensors are never + modified. + + Preconditions, not checked: every id in ``[0, 256)`` and the eight ids of each row + distinct. Mixing weights are used as given; they need not sum to one, and a zero + weight still calls its expert. + """ + _check_inputs(hidden_states, expert_ids, mixing_weights, expert_provider) + return beam_dispatch_combine_op( + hidden_states, + expert_ids, + mixing_weights, + functools.partial(_checked_expert_output, expert_provider), + routed_scaling_factor=ROUTED_SCALING_FACTOR, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/__init__.py new file mode 100644 index 000000000000..23e614d4422b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this component. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/moe/beam_dispatch_combine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/moe/beam_dispatch_combine.py new file mode 100644 index 000000000000..77a4608e7436 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/moe/beam_dispatch_combine.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dispatch tokens to their selected experts and combine the weighted outputs in fp32.""" + +from __future__ import annotations + +from typing import Callable + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Block 3's MoE: 5120 hidden, 256 experts, 8 selected per token, routed scaling factor 2.5. +_BLOCK3 = dict(hidden_size=5120, num_experts=256, top_k=8, routed_scaling_factor=2.5) + + +class BeamDispatchCombine(OpWrapper): + """Run every selected expert on its tokens and sum the weighted outputs per token. + + `hidden_states [T, H]` bf16, `expert_ids [T, K]` int64 and `mixing_weights [T, K]` fp32 + are one routing call's tokens and result: slot `j` of row `t` sends token `t` to expert + `expert_ids[t, j]` with weight `mixing_weights[t, j]`. `expert_provider(e, rows)` is the + expert computation; it returns expert `e`'s unweighted bf16 output, one row per row of + `rows`, in their order. The steps, as the source model computes them:: + + assignments = the T*K slots in token-major order, each with its token and weight + for e in the selected experts, ascending: + tokens = tokens of the assignments to e, ascending [n] int64 + out = expert_provider(e, hidden_states[tokens]) [n, H] bf16 + route = (out.float() * weights[:, None]).to(bfloat16) [n, H] each route rounded once + acc[tokens] += route.float() [T, H] fp32 + output = (acc * routed_scaling_factor).to(bfloat16) [T, H] bf16 + + The provider is called once per selected expert, with at least one row; an expert no + token selected is never called. Returns `output`. `expert_ids`, `mixing_weights` and + `hidden_states` are read, never written; the weights are used as given, without + re-normalization or bias. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block 3, 1 token: eight experts of one row each", spec=dict(_BLOCK3, tokens=1)), + Cell( + why="block 3, 17 tokens: experts shared by several tokens", + spec=dict(_BLOCK3, tokens=17), + ), + Cell( + why="block 3, 129 tokens: almost every expert selected, the longest context", + spec=dict(_BLOCK3, tokens=129), + ), + ) + + note = """ + Every weighted route is rounded to bf16 before it is added, and the sum is fp32. On + finite values that is a different answer from an fp32 route, a bf16 product, a bf16 + accumulator or a sum rounded only at the end, and from folding the scaling factor into + the weights: each moves outputs by bf16 units. A token's eight routes are added in + ascending expert order, one expert at a time; within an expert a token appears once, + so no addition order is left to the hardware. Adding them in slot order, or all at + once with one scattered add, changes the fp32 sum on elements where routes cancel. + + The selected experts and their tokens are host metadata: `unique()` and `tolist()` + synchronize with the device twice for the expert list, and the two boolean selections + twice per selected expert -- 512 synchronizations when 255 experts are selected. Token + rows, weights, products, roundings, sums and the scaling run on the device of + `hidden_states`. + + Not the fused MoE runners of the catalog: those own the expert weights and activation + (and their quantization) as well as the combination, so they cannot call an expert + given as a function, and they combine without the per-route bf16 rounding. + """ + + def __call__( + self, + hidden_states: torch.Tensor, + expert_ids: torch.Tensor, + mixing_weights: torch.Tensor, + expert_provider: Callable[[int, torch.Tensor], torch.Tensor], + routed_scaling_factor: float, + ) -> torch.Tensor: + tokens, top_k = expert_ids.shape + flat_ids = expert_ids.reshape(-1) + flat_weights = mixing_weights.reshape(-1) + token_ids = torch.arange(tokens, device=hidden_states.device).repeat_interleave(top_k) + acc = torch.zeros( + tokens, hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device + ) + for expert_id in flat_ids.unique().tolist(): + assigned = flat_ids == expert_id + expert_tokens = token_ids[assigned] + out = expert_provider(expert_id, hidden_states[expert_tokens]) + route = (out.float() * flat_weights[assigned].unsqueeze(-1)).to(torch.bfloat16) + acc.index_add_(0, expert_tokens, route.float()) + return (acc * routed_scaling_factor).to(torch.bfloat16) + + def reference( + self, + hidden_states: torch.Tensor, + expert_ids: torch.Tensor, + mixing_weights: torch.Tensor, + expert_provider: Callable[[int, torch.Tensor], torch.Tensor], + routed_scaling_factor: float, + ) -> torch.Tensor: + """The combination in plain torch, expert by expert over every id up to the largest. + + Each expert's (token, slot) pairs are looked up in the routing table directly, so + nothing is shared with `__call__`'s flattening and grouping. The product of a bf16 + output and an fp32 weight is exact in fp64; rounding it to fp32 is the fp32 product, + and that is rounded to bf16 as the source does. The bf16 routes are then added in + fp32 in ascending expert order, row by row, and the sum is scaled and rounded once. + Rounding each route is the source's semantics, not a loss of reference accuracy. + """ + acc = torch.zeros( + hidden_states.shape[0], + hidden_states.shape[1], + dtype=torch.float32, + device=hidden_states.device, + ) + for expert_id in range(int(expert_ids.max()) + 1): + token_slot = (expert_ids == expert_id).nonzero() + if token_slot.shape[0] == 0: + continue + token, slot = token_slot[:, 0], token_slot[:, 1] + out = expert_provider(expert_id, hidden_states[token]) + weight = mixing_weights[token, slot].to(torch.float64) + route = (out.to(torch.float64) * weight[:, None]).to(torch.float32).to(torch.bfloat16) + acc[token] = acc[token] + route.to(torch.float32) + return (acc * routed_scaling_factor).to(torch.bfloat16) + + +beam_dispatch_combine_op = BeamDispatchCombine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/__init__.py new file mode 100644 index 000000000000..434f263710d2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One Beam MoE routed expert of decoder block 3 (SwiGLU, H5120 / I2560) / sm_103 / tp1.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/expert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/expert.py new file mode 100644 index 000000000000..d97416018e0f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/expert.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One Beam MoE routed expert of decoder block 3 / sm_103 / tp1, constructed directly. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine builds this class, and it serves no request. +A caller constructs it, loads one expert's three tensors and passes it the +already-normalized hidden states of the tokens routed to that expert. Which expert it is +decides only which tensors are loaded. It routes no token and applies no mixing weight, +routed scaling factor, residual or normalization, and it makes no collective: weighting +and combining expert outputs, and any reduction across ranks, belong to the caller. It +does not call the built-in Beam modeling. + +The expert, bias-free, with the projections in bf16 and the activation in fp32:: + + gate = hidden @ gate_proj^T [T, 2560] bf16 catalog cublas_mm + up = hidden @ up_proj^T [T, 2560] bf16 catalog cublas_mm + silu = silu(gate.float()) [T, 2560] fp32 ops/activation/swiglu_fp32.py + product = silu * up.float() [T, 2560] fp32 + act = product.to(bfloat16) [T, 2560] bf16 the one rounding before down + out = act @ down_proj^T [T, 5120] bf16 catalog cublas_mm +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import Mapping + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm + +from .ops.activation.swiglu_fp32 import swiglu_fp32 + +HIDDEN_SIZE = 5120 +INTERMEDIATE_SIZE = 2560 + +#: The expert's three tensors by checkpoint name (relative to ``mlp.experts..``), [out, in]. +WEIGHT_SHAPES = OrderedDict( + [ + ("gate_proj.weight", (INTERMEDIATE_SIZE, HIDDEN_SIZE)), + ("up_proj.weight", (INTERMEDIATE_SIZE, HIDDEN_SIZE)), + ("down_proj.weight", (HIDDEN_SIZE, INTERMEDIATE_SIZE)), + ] +) + +_SM = (10, 3) + + +def _attribute(name: str) -> str: + return name.removesuffix(".weight") + + +class BeamMoeExpertBf16Sm103Tp1(nn.Module): + """One routed expert: BF16 hidden states in, the expert's unweighted BF16 output out. + + State is exactly three bias-free BF16 matrices, loaded with ``load_weights``: + ``gate_proj`` and ``up_proj`` ``[2560, 5120]`` and ``down_proj`` ``[5120, 2560]``, stored + as given, [out, in]. + + ``forward(hidden_states)`` takes ``[T, 5120]`` ``torch.bfloat16`` on the expert's + device, T >= 1 tokens already normalized by the caller, and returns ``[T, 5120]`` + ``torch.bfloat16`` on that device, one row per input row. Inputs outside that are + rejected before any kernel runs. Measured at T = 1, 17 and 129; zero tokens (an expert + no token was routed to) is rejected, not computed. + + Op ownership: the three projections are the catalog's ``cublas_mm``, each with the + stored weight's ``.t()`` as its column-major operand; the fp32 SiLU, the fp32 product + and the bf16 rounding are this component's ``swiglu_fp32``. + """ + + def __init__(self, device: str | torch.device = "cuda"): + super().__init__() + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this expert is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + for name, shape in WEIGHT_SHAPES.items(): + tensor = torch.empty(shape, dtype=torch.bfloat16, device=device) + self.register_parameter(_attribute(name), nn.Parameter(tensor, requires_grad=False)) + self._loaded = False + + def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": + """The expert's tensors by checkpoint name, in checkpoint order.""" + return OrderedDict((name, getattr(self, _attribute(name))) for name in WEIGHT_SHAPES) + + @torch.no_grad() + def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Copy exactly the expert's three tensors, by checkpoint name: BF16, [out, in]. + + Values are copied as given, without transposing or converting; host tensors are + copied to the expert's device. Every tensor is checked, and copied to the device, + before any parameter changes, so a rejected load leaves the expert as it was. + """ + if not isinstance(weights, Mapping): + raise TypeError(f"weights must be a mapping, got {type(weights).__name__}") + missing = [n for n in WEIGHT_SHAPES if n not in weights] + unexpected = sorted(str(n) for n in weights if n not in WEIGHT_SHAPES) + if missing or unexpected: + raise ValueError(f"expert weights mismatch: missing {missing}, unexpected {unexpected}") + for name, shape in WEIGHT_SHAPES.items(): + tensor = weights[name] + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + if tensor.dtype != torch.bfloat16: + raise TypeError(f"{name} must be torch.bfloat16, got {tensor.dtype}") + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must be {list(shape)}, got {list(tensor.shape)}") + if tensor.layout != torch.strided or tensor.device.type not in ("cpu", "cuda"): + raise ValueError( + f"{name} must be a dense host or CUDA tensor, got {tensor.layout} " + f"on {tensor.device}" + ) + device = self.gate_proj.device + staged = {name: weights[name].to(device=device, copy=True) for name in WEIGHT_SHAPES} + for name, tensor in staged.items(): + getattr(self, _attribute(name)).copy_(tensor) + self._loaded = True + + def _check_input(self, hidden_states: torch.Tensor) -> None: + """Raise on anything this expert would otherwise compute wrongly.""" + if not self._loaded: + raise RuntimeError("load_weights must be called before the expert runs") + if not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"hidden_states must be a torch.Tensor, got {type(hidden_states).__name__}" + ) + if hidden_states.dtype != torch.bfloat16: + raise TypeError(f"hidden_states must be torch.bfloat16, got {hidden_states.dtype}") + if hidden_states.dim() != 2 or hidden_states.shape[1] != HIDDEN_SIZE: + raise ValueError( + f"hidden_states must be [tokens, {HIDDEN_SIZE}], got {list(hidden_states.shape)}" + ) + if hidden_states.shape[0] < 1: + raise ValueError("hidden_states must hold at least one token") + if hidden_states.device != self.gate_proj.device: + raise ValueError( + f"hidden_states must be on the expert's device {self.gate_proj.device}, " + f"got {hidden_states.device}" + ) + + @torch.no_grad() + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self._check_input(hidden_states) + x = hidden_states.contiguous() + gate = cublas_mm(x, self.gate_proj.t()) + up = cublas_mm(x, self.up_proj.t()) + act = swiglu_fp32(gate, up)["act"] + return cublas_mm(act, self.down_proj.t()) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/__init__.py new file mode 100644 index 000000000000..23e614d4422b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this component. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/activation/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/activation/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/activation/swiglu_fp32.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/activation/swiglu_fp32.py new file mode 100644 index 000000000000..b5c482e08a1c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/activation/swiglu_fp32.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SwiGLU with the SiLU and the product in fp32, rounded once to the projections' dtype.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Block 3's experts: the gate and up projections are [T, 2560] bf16. +_EXPERT = dict(intermediate=2560, dtype=torch.bfloat16) + + +class SwigluFp32(OpWrapper): + """`silu(gate.float()) * up.float()`, rounded once to `gate.dtype`, every step returned. + + `gate` and `up` are an expert's two projections, `[T, I]` bf16 each. The steps, as the + source model computes them:: + + silu = silu(gate.float()) [T, I] fp32 + product = silu * up.float() [T, I] fp32 + act = product.to(gate.dtype) [T, I] bf16, the down projection's input + + Returns all three in a dict under those names, so a caller that needs only `act` and a + test that checks each step observe the same call. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block 3 expert, 1 token", spec=dict(_EXPERT, tokens=1)), + Cell(why="block 3 expert, 17 tokens", spec=dict(_EXPERT, tokens=17)), + Cell(why="block 3 expert, 129 tokens: the longest context", spec=dict(_EXPERT, tokens=129)), + ) + + note = """ + Not the shared `flashinfer_silu_and_mul` entry. That kernel also computes in fp32 and + rounds once, but it returns only the rounded activation: the fp32 SiLU and product it + computed cannot be observed, and here they are outputs of the op. + + Rounding the SiLU or the product to bf16 before the final cast is a different answer: + it rounds twice, and on inputs away from any bf16 rounding boundary the activation + moves by a bf16 unit wherever the two roundings disagree. + """ + + def __call__(self, gate: torch.Tensor, up: torch.Tensor) -> dict[str, torch.Tensor]: + silu = F.silu(gate.float()) + product = silu * up.float() + return dict(silu=silu, product=product, act=product.to(gate.dtype)) + + def reference(self, gate: torch.Tensor, up: torch.Tensor) -> dict[str, torch.Tensor]: + """Each step in fp64, rounded to fp32 where the source holds fp32, then to `gate.dtype`. + + The product of an fp32 SiLU and a bf16 `up` is exact in fp64, so `product` is the + correctly rounded fp32 product of the reference's own SiLU. + """ + silu = F.silu(gate.double()).float() + product = (silu.double() * up.double()).float() + return dict(silu=silu, product=product, act=product.to(gate.dtype)) + + def is_valid(self, gate: torch.Tensor, up: torch.Tensor) -> None: + # Mismatched shapes broadcast, and a gate that is not bf16 removes (fp32) or moves + # (fp16) the one rounding the down projection's input carries. Neither raises. + assert gate.shape == up.shape, ( + f"gate {tuple(gate.shape)} and up {tuple(up.shape)} must have one shape" + ) + assert gate.dtype == torch.bfloat16 and up.dtype == torch.bfloat16, ( + f"gate and up must be torch.bfloat16, got {gate.dtype} and {up.dtype}" + ) + + +swiglu_fp32 = SwigluFp32() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/__init__.py new file mode 100644 index 000000000000..539171ed4825 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE router of decoder block 3 (grouped sigmoid top-k) / sm_103 / tp1.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/__init__.py new file mode 100644 index 000000000000..23e614d4422b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this component. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/moe/beam_grouped_topk_router.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/moe/beam_grouped_topk_router.py new file mode 100644 index 000000000000..5a3245af75de --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/moe/beam_grouped_topk_router.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Grouped sigmoid top-k routing with a selection-only correction bias, all in fp32.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: Block 3's router: 256 experts in 8 groups of 32, 4 groups kept, 8 experts per token. +_BLOCK3 = dict(hidden_size=5120, num_experts=256, n_group=8, topk_group=4, top_k=8, eps=1e-20) + + +class BeamGroupedTopkRouter(OpWrapper): + """Route `hidden_states [T, hidden]` with `weight [E, hidden]` and `correction_bias [E]`. + + Every step is fp32, with TF32 and autocast off for the call:: + + logits = hidden_states @ weight^T [T, E] + scores = sigmoid(logits) unbiased + selection_scores = scores + correction_bias used only to choose + group_scores = sum of the two largest selection_scores in each of n_group + contiguous groups [T, n_group] + groups = top-topk_group group_scores, descending [T, topk_group] + masked_selection_scores = selection_scores, -inf outside those groups + ids = top-top_k masked_selection_scores, descending [T, top_k] int64 + weights = scores[ids] / (sum(scores[ids]) + eps) [T, top_k] + + Returns every tensor above in a dict under those names, so a caller that needs only + `ids` and `weights` and a test that checks each step observe the same call. The + correction bias decides the selection and never reaches `weights`; no routed scaling + factor is applied. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block 3, 1 token", spec=dict(_BLOCK3, tokens=1)), + Cell(why="block 3, 17 tokens", spec=dict(_BLOCK3, tokens=17)), + Cell( + why="block 3, 129 tokens: several hundred decision margins", + spec=dict(_BLOCK3, tokens=129), + ), + ) + + note = """ + Ties follow `torch.topk(..., sorted=True)` and nothing else: there is no secondary + index rule. Measured on sm_103 at these shapes, a tie straddling the 4th/5th group or + the 8th/9th expert boundary keeps the lower index, and tied values inside the + selection are not in index order -- a row of 256 equal scores selects groups + [1, 0, 2, 3] and experts [7, 6, 4, 5, 1, 0, 2, 3]. That is the source model's order, + which is why the selection is not delegated to a kernel with its own tie rule. + + The fused grouped top-k op (`noaux_tc_op`) is not this computation: it takes logits, + uses a tanh-form sigmoid that saturates where this one does not, renormalizes in + fp64, applies the routed scaling factor itself and breaks ties by lower expert index. + `reference` keeps the source's `topk` selections because no other formulation + reproduces their order on ties. + """ + + def __call__( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + correction_bias: torch.Tensor, + n_group: int, + topk_group: int, + top_k: int, + eps: float, + ) -> dict[str, torch.Tensor]: + with torch.autocast("cuda", enabled=False): + with true_fp32_matmul(): + logits = F.linear(hidden_states.float(), weight) + scores = logits.sigmoid() + selection_scores = scores + correction_bias + tokens, num_experts = selection_scores.shape + per_group = num_experts // n_group + grouped = selection_scores.view(tokens, n_group, per_group) + group_scores = grouped.topk(2, dim=-1).values.sum(dim=-1) + groups = group_scores.topk(topk_group, dim=-1, sorted=True).indices + group_mask = torch.zeros_like(group_scores).scatter_(1, groups, 1) + expert_mask = ( + group_mask.unsqueeze(-1) + .expand(tokens, n_group, per_group) + .reshape(tokens, num_experts) + ) + masked_selection_scores = selection_scores.masked_fill( + ~expert_mask.bool(), float("-inf") + ) + ids = masked_selection_scores.topk(top_k, dim=-1, sorted=True).indices + selected = scores.gather(-1, ids) + weights = selected / (selected.sum(dim=-1, keepdim=True) + eps) + return dict( + logits=logits, + scores=scores, + selection_scores=selection_scores, + group_scores=group_scores, + groups=groups, + masked_selection_scores=masked_selection_scores, + ids=ids, + weights=weights, + ) + + def reference( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + correction_bias: torch.Tensor, + n_group: int, + topk_group: int, + top_k: int, + eps: float, + ) -> dict[str, torch.Tensor]: + """The routing in plain torch: fp32 operands, TF32 and autocast off. + + Exclusion is written as group membership and each group's score as the first two + entries of a descending sort; both are exact, so the result is the source routing's + bit for bit. The two selections stay `topk(..., sorted=True)`, the source's own + call, because the order of tied scores is whatever that call returns. + """ + with torch.autocast("cuda", enabled=False), true_fp32_matmul(): + logits = F.linear(hidden_states.to(torch.float32), weight.to(torch.float32)) + scores = torch.sigmoid(logits) + selection_scores = scores + correction_bias.to(torch.float32) + tokens, num_experts = selection_scores.shape + per_group = num_experts // n_group + ranked = selection_scores.reshape(tokens, n_group, per_group).sort( + dim=-1, descending=True + ) + group_scores = ranked.values[..., :2].sum(dim=-1) + groups = torch.topk(group_scores, topk_group, dim=-1, sorted=True).indices + expert_group = torch.arange(num_experts, device=logits.device) // per_group + eligible = (expert_group.view(1, num_experts, 1) == groups.unsqueeze(1)).any(dim=-1) + masked_selection_scores = torch.where( + eligible, selection_scores, torch.full_like(selection_scores, float("-inf")) + ) + ids = torch.topk(masked_selection_scores, top_k, dim=-1, sorted=True).indices + selected = torch.take_along_dim(scores, ids, dim=-1) + weights = selected / (selected.sum(dim=-1, keepdim=True) + eps) + return dict( + logits=logits, + scores=scores, + selection_scores=selection_scores, + group_scores=group_scores, + groups=groups, + masked_selection_scores=masked_selection_scores, + ids=ids, + weights=weights, + ) + + def is_valid( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + correction_bias: torch.Tensor, + n_group: int, + topk_group: int, + top_k: int, + eps: float, + ) -> None: + # A [1, T, hidden] batch projects fine, and the group view then folds its leading + # dimension away -- silently when T == 1. + assert hidden_states.dim() == 2, ( + f"hidden_states must be [tokens, hidden], got {tuple(hidden_states.shape)}" + ) + # A length-1 bias broadcasts over every expert, and a wider one promotes every + # selection score out of fp32; neither raises. + assert correction_bias.shape == (weight.shape[0],), ( + f"correction_bias must be [{weight.shape[0]}], got {tuple(correction_bias.shape)}" + ) + assert correction_bias.dtype == torch.float32, ( + f"correction_bias must be torch.float32, got {correction_bias.dtype}" + ) + num_experts = weight.shape[0] + assert num_experts % n_group == 0 and num_experts // n_group >= 2, ( + f"{num_experts} experts do not form {n_group} groups of at least two" + ) + # With fewer eligible experts than top_k, the expert top-k returns excluded ones. + assert 0 < topk_group <= n_group and 0 < top_k <= topk_group * (num_experts // n_group), ( + f"top_k={top_k} cannot be chosen from {topk_group} of {n_group} groups" + ) + + +beam_grouped_topk_router = BeamGroupedTopkRouter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/router.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/router.py new file mode 100644 index 000000000000..8a88d7d0309f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/router.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE router of the sparse decoder blocks 3-5 / sm_103 / tp1, constructed directly. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine builds this class, and it serves no request. +A caller constructs it from the checkpoint's configuration, loads the router's two +tensors and passes it the already-normalized hidden states that enter its sparse +block's feed-forward sublayer; each sparse block owns its own router. It returns the +eight experts each token is routed to and their mixing weights. It runs no expert, +dispatches no token and combines no expert output; the routed scaling factor belongs to +that combine and is not applied here. It does not call the built-in Beam modeling. + +The routing, every step fp32 (``ops/moe/beam_grouped_topk_router.py``):: + + logits = hidden.float() @ weight^T [T, 256] + scores = sigmoid(logits) unbiased: the weights use these + choice = scores + e_score_correction_bias biased: only selection uses these + groups = top-4 of (sum of the 2 largest choice in each of 8 groups of 32), descending + ids = top-8 choice among the experts of those groups, descending [T, 8] int64 + weights = scores[ids] / (sum(scores[ids]) + 1e-20) [T, 8] fp32 +""" + +from __future__ import annotations + +import torch +from torch import nn + +from .ops.moe.beam_grouped_topk_router import beam_grouped_topk_router + +#: The configuration this component computes, as the checkpoint's config.json states it. +#: Anything else is rejected at construction rather than computed approximately. +SUPPORTED_CONFIG = { + "hidden_size": 5120, + "num_experts": 256, + "num_experts_per_tok": 8, + "n_group": 8, + "topk_group": 4, + "norm_topk_prob": True, +} +SUPPORTED_LAYER_IDX = 3 +#: The sparse blocks this router has been built for: the checkpoint's first three. +SUPPORTED_LAYER_INDICES = (3, 4, 5) + +#: Added to the selected scores' sum before dividing, as the source model does. +NORM_EPS = 1e-20 + +_SM = (10, 3) + + +def check_config(config, layer_idx: int) -> None: + """Raise ValueError naming the first field outside what this component computes.""" + if type(layer_idx) is not int or layer_idx not in SUPPORTED_LAYER_INDICES: + raise ValueError( + f"only the routers of sparse blocks {SUPPORTED_LAYER_INDICES} are supported, " + f"got layer_idx={layer_idx!r}" + ) + for field, expected in SUPPORTED_CONFIG.items(): + actual = getattr(config, field, None) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this router computes {expected!r}") + mlp_types = getattr(config, "mlp_layer_types", None) or [] + if len(mlp_types) <= layer_idx or mlp_types[layer_idx] != "sparse": + raise ValueError( + f"block {layer_idx} must have a sparse (MoE) feed-forward: " + f"mlp_layer_types[{layer_idx}]=" + f"{mlp_types[layer_idx] if len(mlp_types) > layer_idx else None!r}" + ) + + +class BeamMoeRouterBf16Sm103Tp1(nn.Module): + """A sparse block's router: BF16 hidden states in, expert ids and FP32 mixing weights out. + + State is exactly two FP32 tensors, loaded with ``load_weights``: ``weight`` + ``[256, 5120]`` (checkpoint ``mlp.gate.weight``) and ``e_score_correction_bias`` + ``[256]`` (``mlp.gate.e_score_correction_bias``). + + ``forward(hidden_states)`` takes ``[T, 5120]`` ``torch.bfloat16`` on the router's + device, T >= 1 tokens already normalized by the caller, and returns + ``(mixing_weights, expert_ids)`` -- the source model's order -- both ``[T, 8]`` on that + device: ``mixing_weights`` ``torch.float32``, ``expert_ids`` ``torch.int64`` in + ``[0, 256)``, in descending order of the biased selection score. Inputs outside that + are rejected before any kernel runs. Routing runs entirely on the GPU and makes no + host synchronization. + """ + + def __init__( + self, config, layer_idx: int = SUPPORTED_LAYER_IDX, device: str | torch.device = "cuda" + ): + super().__init__() + check_config(config, layer_idx) + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this router is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_tok + self.n_group = config.n_group + self.topk_group = config.topk_group + self.weight = nn.Parameter( + torch.empty(self.num_experts, self.hidden_size, dtype=torch.float32, device=device), + requires_grad=False, + ) + self.register_buffer( + "e_score_correction_bias", + torch.empty(self.num_experts, dtype=torch.float32, device=device), + ) + self._loaded = False + + @torch.no_grad() + def load_weights(self, weight: torch.Tensor, correction_bias: torch.Tensor) -> None: + """Copy the router's two tensors: ``weight [256, 5120]`` and ``correction_bias [256]``. + + Both must already be ``torch.float32``: the router computes in fp32 and does not + accept parameters that were rounded below it. Both are checked before either is + copied, so a rejected load leaves the router as it was. + """ + expected = ( + ("weight", weight, (self.num_experts, self.hidden_size)), + ("correction_bias", correction_bias, (self.num_experts,)), + ) + for name, tensor, shape in expected: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + if tensor.dtype != torch.float32: + raise TypeError(f"{name} must be torch.float32, got {tensor.dtype}") + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must be {list(shape)}, got {list(tensor.shape)}") + self.weight.copy_(weight) + self.e_score_correction_bias.copy_(correction_bias) + self._loaded = True + + def _check_input(self, hidden_states: torch.Tensor) -> None: + """Raise on anything this router would otherwise compute wrongly or move to the host.""" + if not self._loaded: + raise RuntimeError("load_weights must be called before routing") + if not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"hidden_states must be a torch.Tensor, got {type(hidden_states).__name__}" + ) + if hidden_states.dtype != torch.bfloat16: + raise TypeError(f"hidden_states must be torch.bfloat16, got {hidden_states.dtype}") + if hidden_states.dim() != 2 or hidden_states.shape[1] != self.hidden_size: + raise ValueError( + f"hidden_states must be [tokens, {self.hidden_size}], " + f"got {list(hidden_states.shape)}" + ) + if hidden_states.shape[0] < 1: + raise ValueError("hidden_states must hold at least one token") + if hidden_states.device != self.weight.device: + raise ValueError( + f"hidden_states must be on the router's device {self.weight.device}, " + f"got {hidden_states.device}" + ) + + def route(self, hidden_states: torch.Tensor) -> dict[str, torch.Tensor]: + """Every intermediate of one routing call, by the op's names; see the op's docstring.""" + self._check_input(hidden_states) + return beam_grouped_topk_router( + hidden_states, + self.weight, + self.e_score_correction_bias, + n_group=self.n_group, + topk_group=self.topk_group, + top_k=self.top_k, + eps=NORM_EPS, + ) + + @torch.no_grad() + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + routing = self.route(hidden_states) + return routing["weights"], routing["ids"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/routing.py new file mode 100644 index 000000000000..9c8a301a870d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/routing.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam family routing: which ModelingV2 target serves a ``BeamForCausalLM`` deployment.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Optional, Tuple + +from ..._router_index import NoTarget + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + +#: Proposals per DFlash draft block the target verifies: the drafter's serving block of eight +#: rows (the pending token and seven proposals), the largest verification width the target's +#: attention op is measured at. +DFLASH_MAX_DRAFT_LEN = 7 + + +def _quant_algo(pretrained_config: Any) -> Optional[str]: + quant = getattr(pretrained_config, "quantization_config", None) + if isinstance(quant, dict): + return quant.get("quant_algo") + return getattr(quant, "quant_algo", None) + + +def speculation_refusal(spec_config: Any, llm_args: Optional[Any] = None) -> Optional[str]: + """Why the Beam target cannot run ``spec_config``, or None when it can. + + The one speculative mode is DFlash with a separate drafter checkpoint, the TRTLLM draft + attention backend (the drafter's context in the engine's paged draft KV pool), a fixed + draft length of ``DFLASH_MAX_DRAFT_LEN`` and strict greedy verification (no rejection + sampling, no draft-length schedule). The engine must run it eagerly: the overlap + scheduler would hand the target device-corrected positions and KV lengths its host-side + step check cannot state. + """ + if spec_config is None: + return None + mode = getattr(spec_config, "spec_dec_mode", None) + if mode is None or not mode.is_dflash(): + return f"speculative decoding {type(spec_config).__name__} is not DFlash" + if spec_config.speculative_model is None: + return "DFlash needs its drafter checkpoint (speculative_model)" + if spec_config.attention_backend != "TRTLLM": + return ( + f"DFlash attention_backend {spec_config.attention_backend!r}; the drafter's context " + "lives in the paged draft KV pool, which only the TRTLLM backend reads" + ) + if spec_config.max_draft_len != DFLASH_MAX_DRAFT_LEN: + return f"DFlash max_draft_len {spec_config.max_draft_len}, not {DFLASH_MAX_DRAFT_LEN}" + if spec_config.draft_len_schedule: + return "a DFlash draft_len_schedule changes the verification width per step" + if spec_config.use_rejection_sampling: + return "DFlash rejection sampling; verification is strict greedy" + if llm_args is not None and not llm_args.disable_overlap_scheduler: + return "the overlap scheduler with DFlash" + return None + + +def route( + pretrained_config: Any, + mapping: Any, + compute_capability: Tuple[int, int], + spec_config: Optional[Any], + kv_cache_quant_algo: Optional[Any], + model_config: Optional[ModelConfig] = None, + llm_args: Optional[TorchLlmArgs] = None, +) -> type: + """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. + + The one target is the 52-layer NVFP4-expert checkpoint on four sm_103 ranks: attention + TP4, MoE TP4/EP1, PP1, BF16 KV cache, no attention data parallelism, without speculative + decoding or with the DFlash drafter ``speculation_refusal`` admits. ``kv_cache_quant_algo`` + is the effective engine setting, including runtime overrides. No other engine choice in + ``model_config`` or ``llm_args`` selects a different target. + """ + shape = ( + pretrained_config.num_hidden_layers, + pretrained_config.hidden_size, + getattr(pretrained_config, "num_experts", None), + ) + layout = ( + mapping.world_size, + mapping.tp_size, + mapping.pp_size, + mapping.moe_tp_size, + mapping.moe_ep_size, + bool(mapping.enable_attention_dp), + ) + quant_algo = _quant_algo(pretrained_config) + speculation = speculation_refusal(spec_config, llm_args) + if ( + compute_capability == (10, 3) + and shape == (52, 5120, 256) + and layout == (4, 4, 1, 4, 1, False) + and quant_algo == "NVFP4" + and speculation is None + and kv_cache_quant_algo is None + ): + from .full_model_nvfp4__sm_103__tp4.modeling import ModelingV2BeamFullModelNvfp4Sm103Tp4 + + return ModelingV2BeamFullModelNvfp4Sm103Tp4 + + raise NoTarget( + f"compute capability {compute_capability}, config shape {shape}, quantization " + f"{quant_algo}, layout (world, tp, pp, moe_tp, moe_ep, attention_dp) {layout}, " + f"speculative decoding {type(spec_config).__name__ if spec_config is not None else 'off'}" + f"{f' ({speculation})' if speculation else ''}; " + f"KV cache {kv_cache_quant_algo or 'bf16'}; " + "the one Beam target is the 52-layer NVFP4 checkpoint on sm103, attention TP4, MoE " + "TP4/EP1, PP1, with BF16 KV cache and without attention DP, with speculative decoding " + f"off or DFlash (TRTLLM draft attention, max_draft_len {DFLASH_MAX_DRAFT_LEN}, strict " + "greedy verification, no overlap scheduler)" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/__init__.py new file mode 100644 index 000000000000..77397f5b34ef --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam six-block small model, token ids to prefill logits (H5120, vocab 151936) / sm_103 / tp1.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/generation.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/generation.py new file mode 100644 index 000000000000..17dbf191dde9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/generation.py @@ -0,0 +1,694 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Stateful greedy generation for the Beam six-block model / sm_103 / tp1. + +A thin serial controller that drives the stateless six-block model +(``BeamSmallModelPrefillBf16Sm103Tp1``) through TensorRT-LLM's own request, cache and +sampler objects: one ``LlmRequest`` per request, one long-lived ``KVCacheManagerV2`` with +block reuse for all of them, a ``SeqSlotManager`` for sampler slots and the existing +``TorchSampler``. It owns no cache allocator, prefix table or sampler of its own; it calls +those objects in the order the non-overlap executor does -- forward, ``sample_async``, +context-cursor advance, ``update_requests``, ``update_context_resources``, finish +handling, ``update_resources``. It is a component, not a serving engine: one request +executes at a time, at most ``GENERATOR_CACHE["max_batch_size"]`` (two) request states +are live, and there is no scheduler, batching, streaming or LLM API registration. + +Lifecycle:: + + with BeamSmallModelGenerator(model) as gen: # binds the loaded model + state = gen.create_request(prompt_ids, max_new_tokens=8, end_id=None) + first = gen.prefill(state) # prompt tokens not already cached; first token + while state.lifecycle is not Lifecycle.FINISHED: + step = gen.decode(state) # consumes the pending sampled token + tokens, reason = state.output_token_ids, state.finish_reason + gen.release(state) + +``generate(state)`` runs the prefill and decode steps of a new sampling request to its +end. A compute-only request (``sample=False``) runs the same steps without the sampler: +``prefill`` returns the last prompt token's logits and ``decode(state, token_id)`` +consumes an explicit (teacher-forced) token; its token budget is ``max_new_tokens`` +explicit tokens. + +Semantics. + +* Positions are absolute. A prefill computes prompt tokens ``reused .. P - 1`` at those + positions after ``reused`` cached ones; a decode consumes one token at position + ``consumed`` (the number of tokens whose keys and values the cache holds) and leaves + ``consumed + 1``. Every model call is checked by the model against the manager before + it runs (``model.check_request``). +* Prefix reuse is the manager's. A prefill first asks the manager for the request's + cache (``prepare_context``): the manager matches the prompt against what earlier + requests committed and claims the longest reusable full-page prefix -- at most + ``P - 1`` tokens, since the last prompt token is always recomputed for its logits; + with partial reuse off, only whole 32-token pages. The prefill computes the rest and the + manager commits it (``update_context_resources``), so later requests can reuse it, + also after this request is released. ``reuse=False`` claims nothing + (``prepare_context_cache(reuse_limit=0)``) and computes the whole prompt: the + reuse-disabled control. Generated tokens are not committed for reuse. +* Sampling is greedy: ``temperature=0``, ``top_k=1``, beam width 1, no penalties, stop + words, logits processors or draft tokens. ``TorchSampler`` chooses each token from the + step's FP32 logits (the lowest vocabulary index among exact maxima). +* Stopping follows the sampler: ``FinishReason.END_ID`` when the sampled token equals + ``end_id``, ``FinishReason.LENGTH`` after ``max_new_tokens`` sampled tokens; END_ID wins + when both hold. The last sampled token is returned and never consumed: after P prompt + tokens and G sampled tokens the cache holds P + G - 1 tokens, and no model call follows + the finishing step. An explicit token never counts as sampled. +* Isolation: each request owns its ``LlmRequest``, cache, sampler slot and ledger; the + model holds no request state. Immutable committed prefix pages may be shared between + requests at the same layer; every page a request writes is its own. +* Identity: a generator binds one model and its current weights -- each of the model's + checkpoint tensors as the very tensor object it is, with its version counter, storage + address, shape and strides. While it is open the model refuses ``load_weights``, and + every step first verifies that the model's checkpoint tensors are still exactly those + bound ones, unmodified: a tensor changed in place, replaced by another tensor (for + example ``load_state_dict(..., assign=True)`` or attribute assignment) or given other + storage is refused before the step looks up a prefix, allocates or runs a kernel, for as + long as the change persists. ``close`` releases every request and shuts the manager + down, after which a new generator starts from an empty cache and binds the weights the + model then holds. + +Limits: requests of 1 .. ``max_seq_len - max_new_tokens`` prompt tokens, ``P + +max_new_tokens <= 1024``; at most two live states; one GB300 (sm_103), eager execution only +(no CUDA graph, no overlap scheduler): every step synchronizes the host several times. +Requests are validated before any kernel runs; a request whose step failed after it began +is FAILED and can only be released. +""" + +from __future__ import annotations + +import enum +import itertools +import weakref +from dataclasses import dataclass +from typing import Optional, Sequence + +import numpy as np +import torch + +from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp +from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.pyexecutor.seq_slot_manager import SeqSlotManager +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.executor import FinishReason +from tensorrt_llm.bindings.executor import SamplingConfig as ExecutorSamplingConfig +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping as ParallelMapping + +from .model import ATTENTION_WINDOWS, MAX_TOKENS, NUM_LAYERS, BeamSmallModelPrefillBf16Sm103Tp1 + +#: The generator's cache: bf16, 32-token pages, a 1024-token sequence capacity, two live +#: requests, block reuse of whole pages (partial reuse off), no host tier and no +#: sliding-window scratch pages. ``max_tokens`` sizes the device quota so that two live +#: requests at full capacity fit beside the reusable pages of released ones. +GENERATOR_CACHE = { + "tokens_per_block": 32, + "max_seq_len": 1024, + "max_batch_size": 2, + "max_tokens": 4096, +} + +#: The sampler: greedy, beam width 1, no draft tokens, no overlap scheduler or async worker. +SAMPLER_ARGS = TorchSampler.Args( + max_seq_len=GENERATOR_CACHE["max_seq_len"], + max_draft_len=0, + max_num_sequences=GENERATOR_CACHE["max_batch_size"], + max_beam_width=1, + max_total_draft_tokens=0, + disable_overlap_scheduler=True, + enable_async_worker=False, +) + + +def build_cache_manager(vocab_size: int = 151936) -> KVCacheManagerV2: + """The generator's six-layer ``KVCacheManagerV2`` (``GENERATOR_CACHE``), block reuse on. + + Logical layers 0-5 in slots 0-5, 8 bf16 kv heads of 128, per-layer windows + ``ATTENTION_WINDOWS`` (V2 keeps the five sliding layers and layer 5 in separate pools), + whole-page reuse under the ``all_reusable`` policy, no host tier, no scratch pages, + no overlap scheduler. The caller shuts it down. + """ + cache = GENERATOR_CACHE + return KVCacheManagerV2( + KvCacheConfig( + max_tokens=cache["max_tokens"], + max_attention_window=list(ATTENTION_WINDOWS), + enable_block_reuse=True, + enable_partial_reuse=False, + enable_swa_scratch_reuse=False, + host_cache_size=0, + ), + CacheTypeCpp.SELF, + num_layers=NUM_LAYERS, + num_kv_heads=8, + head_dim=128, + tokens_per_block=cache["tokens_per_block"], + max_seq_len=cache["max_seq_len"], + max_batch_size=cache["max_batch_size"], + mapping=ParallelMapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.BF16, + vocab_size=vocab_size, + disable_overlap_scheduler=True, + ) + + +def _tensor_state(tensor: torch.Tensor) -> tuple: + """What changes when a tensor's contents or storage change: version, address, shape, strides.""" + return tensor._version, tensor.data_ptr(), tuple(tensor.shape), tensor.stride() + + +def greedy_sampling_config() -> ExecutorSamplingConfig: + """The request sampling config: temperature 0 and top_k 1 at beam width 1, nothing else set. + + ``TorchSampler`` resolves it to ``SamplingMode.GREEDY``. + """ + return ExecutorSamplingConfig(beam_width=1, top_k=1, temperature=0.0) + + +def sample_step( + sampler: TorchSampler, request: LlmRequest, logits: torch.Tensor, *, is_context: bool +) -> None: + """Choose ``request``'s next token from ``logits`` (``[vocab]`` FP32) with ``sampler``. + + Runs ``sample_async`` on a batch holding only ``request`` -- as its last context chunk + or as its one generation row -- then advances a context request's cursor to the end of + its prompt as the executor does, then ``update_requests``, which appends the chosen + token to ``request`` and finishes it when the sampler's stopping rules hold. The + request must already hold a sampler slot. + """ + batch = ScheduledRequests() + if is_context: + batch.context_requests_last_chunk = [request] + prefix_sum = [0, 1] + else: + batch.generation_requests = [request] + prefix_sum = [0] + sample_state = sampler.sample_async(batch, {"logits": logits.view(1, -1)}, prefix_sum) + if is_context: + request.py_last_context_chunk = ( + request.context_current_position, + request.context_current_position + request.context_chunk_size, + ) + request.move_to_next_context_chunk() + if request.context_remaining_length == 0: + request.state = LlmRequestState.GENERATION_IN_PROGRESS + sampler.update_requests(sample_state) + + +class _Request(LlmRequest): + """An ``LlmRequest`` that keeps the finish reason the sampler records on it.""" + + py_finish_reason: Optional[FinishReason] = None + + def finish_by(self, reason: FinishReason, beam: int) -> None: + super().finish_by(reason, beam) + self.py_finish_reason = reason + + +class Lifecycle(enum.Enum): + NEW = "new" + PREFILLED = "prefilled" + DECODING = "decoding" + FINISHED = "finished" + RELEASED = "released" + FAILED = "failed" + + +@dataclass(frozen=True) +class StepOutput: + """One model step of one request. + + ``logits`` are the FP32 ``[151936]`` softcapped logits of the step's last token, at + absolute ``position``. ``token_id`` is the token the sampler chose from them (None for + a compute-only request) and ``finish_reason`` is set on the step that finished it. + """ + + logits: torch.Tensor + position: int + num_computed: int + token_id: Optional[int] + finish_reason: Optional[FinishReason] + + +@dataclass(frozen=True) +class GenerationOutput: + token_ids: tuple + finish_reason: FinishReason + steps: tuple + + +class RequestState: + """A handle to one request of one generator; read-only to callers. + + ``consumed``: tokens whose keys and values the cache holds. ``reused``: prompt tokens + the manager claimed from earlier requests. ``computed_prompt_tokens``: prompt tokens + the prefill computed. ``output_token_ids``: sampled tokens, in order. + ``ledger``: one record per step (positions, token source, manager claim, capacity, + history, committed tokens and page indices after the step). + """ + + def __init__( + self, + generator, + request: LlmRequest, + prompt: tuple, + *, + sample: bool, + reuse: bool, + max_new_tokens: int, + end_id: Optional[int], + ): + self._generator = weakref.ref(generator) + self._request = request + self.request_id = request.py_request_id + self.prompt = prompt + self.sample = sample + self.reuse = reuse + self.max_new_tokens = max_new_tokens + self.end_id = end_id + self.lifecycle = Lifecycle.NEW + self.consumed = 0 + self.reused = 0 + self.computed_prompt_tokens = 0 + self.explicit_tokens = 0 + self.output_token_ids: list = [] + self.pending_token: Optional[int] = None + self.finish_reason: Optional[FinishReason] = None + self.ledger: list = [] + + @property + def prompt_len(self) -> int: + return len(self.prompt) + + def __repr__(self) -> str: + return ( + f"RequestState(id={self.request_id}, {self.lifecycle.value}, consumed={self.consumed}, " + f"reused={self.reused}, output={self.output_token_ids})" + ) + + +class BeamSmallModelGenerator: + """Greedy cached generation and cross-request prefix reuse for one loaded six-block model. + + See the module docstring for the lifecycle, semantics and limits. ``model`` must be a + loaded ``BeamSmallModelPrefillBf16Sm103Tp1``; the generator builds its own + ``KVCacheManagerV2`` (``GENERATOR_CACHE``), ``SeqSlotManager`` and ``TorchSampler`` + (``SAMPLER_ARGS``) and releases them in ``close``. + """ + + _request_ids = itertools.count(1) + + def __init__(self, model: BeamSmallModelPrefillBf16Sm103Tp1): + if not isinstance(model, BeamSmallModelPrefillBf16Sm103Tp1): + raise TypeError( + f"model must be BeamSmallModelPrefillBf16Sm103Tp1, got {type(model).__name__}" + ) + if not model._loaded: + raise ValueError("load_weights must complete before a generator binds the model") + self.model = model + self.closed = False + self._weights = { + name: (tensor, _tensor_state(tensor)) + for name, tensor in model.checkpoint_tensors().items() + } + self.manager = build_cache_manager(model.vocab_size) + try: + self.seq_slots = SeqSlotManager(GENERATOR_CACHE["max_batch_size"]) + self.sampler = TorchSampler(SAMPLER_ARGS) + except BaseException: + self.manager.shutdown() + raise + self._live: dict = {} + model.register_cache_holder(self) + + # ------------------------------------------------------------------ lifecycle + + def __enter__(self) -> "BeamSmallModelGenerator": + return self + + def __exit__(self, *exc) -> None: + self.close() + + @property + def live_requests(self) -> tuple: + return tuple(self._live.values()) + + def create_request( + self, + prompt_ids, + *, + max_new_tokens: int, + end_id: Optional[int] = None, + sample: bool = True, + reuse: bool = True, + ) -> RequestState: + """Admit a new request; nothing runs on the GPU yet. + + ``prompt_ids``: a non-empty sequence (or 1-D integer tensor) of token ids in + ``[0, vocab)``, copied. ``max_new_tokens >= 1``: tokens to sample (a sampling + request) or explicit tokens to consume (``sample=False``), with + ``P + max_new_tokens <= max_seq_len``. ``end_id``: a token id that stops sampling, + or None. ``reuse``: whether the prefill may claim a cached prefix. Raises + ValueError, without side effects, on anything else or when two states are live. + """ + self._check_open() + if isinstance(prompt_ids, torch.Tensor): + if ( + prompt_ids.dim() != 1 + or prompt_ids.is_floating_point() + or prompt_ids.dtype == torch.bool + ): + raise ValueError("prompt_ids must be a 1-D integer tensor") + prompt = tuple(int(t) for t in prompt_ids.tolist()) + else: + prompt = tuple(prompt_ids) + if not all( + isinstance(t, (int, np.integer)) and not isinstance(t, bool) for t in prompt + ): + raise ValueError("prompt_ids must hold integer token ids") + prompt = tuple(int(t) for t in prompt) + vocab = self.model.vocab_size + capacity = self.manager.max_seq_len + if not prompt: + raise ValueError("the prompt must hold at least one token") + if any(not 0 <= t < vocab for t in prompt): + raise ValueError(f"prompt token ids must be in [0, {vocab})") + if type(max_new_tokens) is not int or max_new_tokens < 1: + raise ValueError(f"max_new_tokens must be a positive int, got {max_new_tokens!r}") + if len(prompt) + max_new_tokens > capacity: + raise ValueError( + f"{len(prompt)} prompt tokens + {max_new_tokens} new tokens exceed the " + f"{capacity}-token sequence capacity" + ) + if end_id is not None and (type(end_id) is not int or not 0 <= end_id < vocab): + raise ValueError(f"end_id must be None or a token id in [0, {vocab}), got {end_id!r}") + if end_id is not None and not sample: + raise ValueError("end_id stops sampling; a compute-only request samples nothing") + if len(self._live) >= GENERATOR_CACHE["max_batch_size"]: + raise ValueError( + f"{len(self._live)} requests are live; release one before creating another" + ) + request = _Request( + request_id=next(self._request_ids), + max_new_tokens=max_new_tokens, + input_tokens=list(prompt), + sampling_config=greedy_sampling_config(), + is_streaming=False, + end_id=end_id, + exclude_input_from_output=True, + ) + state = RequestState( + self, + request, + prompt, + sample=sample, + reuse=reuse, + max_new_tokens=max_new_tokens, + end_id=end_id, + ) + self._live[state.request_id] = state + return state + + def release(self, state: RequestState) -> None: + """Free the request's cache and sampler slot; its handle can no longer be used. + + Waits for issued work first. Committed pages stay reusable by later requests until + the manager evicts them. Releasing a released or foreign handle raises ValueError. + """ + self._check_open() + self._check_state(state, tuple(Lifecycle)) + torch.cuda.current_stream().synchronize() + request = state._request + try: + self.manager.free_resources(request) + if request.py_seq_slot is not None: + self.seq_slots.free_resources(request) + finally: + state.lifecycle = Lifecycle.RELEASED + del self._live[state.request_id] + + def close(self) -> None: + """Release every live request and shut the manager down. Idempotent.""" + if self.closed: + return + try: + torch.cuda.current_stream().synchronize() + for state in list(self._live.values()): + try: + self.manager.free_resources(state._request) + finally: + state.lifecycle = Lifecycle.RELEASED + self._live.clear() + finally: + self.closed = True + self.manager.shutdown() + + # ------------------------------------------------------------------ steps + + def prefill(self, state: RequestState) -> StepOutput: + """Compute the prompt tokens the cache does not hold yet, after any reused prefix. + + Claims the reusable prefix through the manager (none with ``reuse=False``), sizes + the cache for the prompt, runs the model on tokens ``reused .. P - 1`` and commits + them for reuse. A sampling request then samples its first token, which may finish + it. Returns the last prompt token's logits. + """ + self._check_open() + self._check_state(state, (Lifecycle.NEW,)) + self._check_weights() + request = state._request + try: + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + self.seq_slots.prepare_resources(batch) + if not state.reuse: + claim = self.manager.prepare_context_cache(request, reuse_limit=0) + cache = self.manager.kv_cache_map.get(state.request_id) + if claim != 0 or cache is None or cache.num_committed_tokens != 0: + raise RuntimeError(f"the reuse-disabled control claimed {claim} tokens") + if not self.manager.prepare_context(request): + raise RuntimeError(f"the manager could not admit request {state.request_id}") + reused = request.context_current_position + num_tokens = request.context_remaining_length + if reused + num_tokens != state.prompt_len or num_tokens < 1: + raise RuntimeError( + f"the manager claimed {reused} of {state.prompt_len} prompt tokens and left " + f"{num_tokens} to compute" + ) + if not self.manager.resize_context(request, num_tokens): + raise RuntimeError(f"no cache capacity for request {state.request_id}'s prompt") + self.manager.prepare_resources(batch) + state.reused = reused + state.computed_prompt_tokens = num_tokens + logits = self._forward(state, state.prompt[reused:], reused, is_context=True) + row = logits[0, -1] + out_logits = row.clone() + token = reason = None + if state.sample: + sample_step(self.sampler, request, row, is_context=True) + else: + request.py_last_context_chunk = (reused, reused + num_tokens) + request.move_to_next_context_chunk() + request.state = LlmRequestState.GENERATION_IN_PROGRESS + self.manager.update_context_resources(batch) + state.consumed = state.prompt_len + if state.sample: + token, reason = self._after_sampling(state) + state.lifecycle = Lifecycle.FINISHED if reason is not None else Lifecycle.PREFILLED + self._record( + state, "prefill", reused, num_tokens, "prompt", token, reason, claimed=reused + ) + return StepOutput(out_logits, state.prompt_len - 1, num_tokens, token, reason) + except Exception: + state.lifecycle = Lifecycle.FAILED + raise + + def decode(self, state: RequestState, token_id: Optional[int] = None) -> StepOutput: + """Consume one token at position ``consumed`` and return the next token's logits. + + A sampling request consumes its pending sampled token (``token_id`` must be None or + equal to it) and samples the next one; a compute-only request consumes the explicit + ``token_id``. The cache grows by one token, which the step writes in every layer. + """ + self._check_open() + self._check_state(state, (Lifecycle.PREFILLED, Lifecycle.DECODING)) + request = state._request + if state.sample: + if state.pending_token is None: + raise ValueError("a sampling request has no pending token to consume") + if token_id is not None and token_id != state.pending_token: + raise ValueError( + f"a sampling request consumes its sampled token {state.pending_token}, " + f"not {token_id}" + ) + token, source = state.pending_token, "sampled" + else: + if ( + token_id is None + or type(token_id) is not int + or not 0 <= token_id < self.model.vocab_size + ): + raise ValueError( + f"a compute-only request needs an explicit token id, got {token_id!r}" + ) + if state.explicit_tokens >= state.max_new_tokens: + raise ValueError(f"the request's budget of {state.max_new_tokens} tokens is spent") + token, source = token_id, "teacher_forced" + if state.consumed + 1 > self.manager.max_seq_len: + raise ValueError(f"the request already holds {state.consumed} tokens, the capacity") + self._check_weights() + try: + if not state.sample: + request.add_new_token(token, 0) + state.explicit_tokens += 1 + past = state.consumed + if request.max_beam_num_tokens - 1 != past or request.get_last_tokens(0) != token: + raise RuntimeError( + f"request {state.request_id} holds {request.max_beam_num_tokens} tokens " + f"ending in {request.get_last_tokens(0)}; the step consumes {token} at {past}" + ) + if not self.manager.try_allocate_generation(request): + raise RuntimeError(f"no cache capacity for request {state.request_id}'s next token") + batch = ScheduledRequests() + batch.generation_requests = [request] + self.manager.prepare_resources(batch) + logits = self._forward(state, (token,), past, is_context=False) + row = logits[0, -1] + out_logits = row.clone() + state.consumed = past + 1 + state.pending_token = None + reason = sampled = None + if state.sample: + sample_step(self.sampler, request, row, is_context=False) + sampled, reason = self._after_sampling(state) + self.manager.update_resources(batch) + else: + self.manager.update_resources(batch) + cache = self.manager.kv_cache_map[state.request_id] + if not cache.resize(None, state.consumed) or cache.history_length != state.consumed: + raise RuntimeError(f"the manager did not record {state.consumed} cached tokens") + state.lifecycle = Lifecycle.FINISHED if reason is not None else Lifecycle.DECODING + self._record(state, "decode", past, 1, source, sampled, reason) + return StepOutput(out_logits, past, 1, sampled, reason) + except Exception: + state.lifecycle = Lifecycle.FAILED + raise + + def generate(self, state: RequestState) -> GenerationOutput: + """Prefill a new sampling request and decode until the sampler finishes it.""" + self._check_open() + self._check_state(state, (Lifecycle.NEW,)) + if not state.sample: + raise ValueError("generate needs a sampling request") + steps = [self.prefill(state)] + while state.lifecycle is not Lifecycle.FINISHED: + steps.append(self.decode(state)) + return GenerationOutput(tuple(state.output_token_ids), state.finish_reason, tuple(steps)) + + # ------------------------------------------------------------------ internals + + def _check_open(self) -> None: + if self.closed: + raise ValueError("the generator is closed") + + def _check_state(self, state, allowed) -> None: + if not isinstance(state, RequestState) or state._generator() is not self: + raise ValueError("the request belongs to another generator") + if self._live.get(state.request_id) is not state: + raise ValueError(f"request {state.request_id} was released") + if state.lifecycle not in allowed: + raise ValueError( + f"request {state.request_id} is {state.lifecycle.value}; this call needs " + f"{', '.join(s.value for s in allowed)}" + ) + + def _check_weights(self) -> None: + """Raise ValueError unless the model's checkpoint tensors are the bound, unmodified ones.""" + current = self.model.checkpoint_tensors() + if current.keys() != self._weights.keys(): + raise ValueError( + "the model's tensors changed after the generator bound it; its cache no longer " + "matches the weights. Close the generator and build a new one" + ) + for name, tensor in current.items(): + bound, state = self._weights[name] + if tensor is not bound or _tensor_state(tensor) != state: + raise ValueError( + f"{name} changed after the generator bound the model; its cache no longer " + "matches the weights. Close the generator and build a new one" + ) + + def _metadata(self, state: RequestState, num_tokens: int, past: int, is_context: bool): + md = TrtllmAttentionMetadata( + max_num_requests=1, + max_num_tokens=MAX_TOKENS, + kv_cache_manager=self.manager, + num_contexts=1 if is_context else 0, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[past]), + seq_lens=torch.tensor([num_tokens], dtype=torch.int), + request_ids=[state.request_id], + prompt_lens=[num_tokens if is_context else state.prompt_len], + runtime_features=AttentionRuntimeFeatures(cache_reuse=True), + ) + md.prepare() + return md + + def _forward(self, state: RequestState, tokens: Sequence[int], past: int, is_context: bool): + num_tokens = len(tokens) + device = self.model.device + md = self._metadata(state, num_tokens, past, is_context) + ids = torch.tensor([list(tokens)], dtype=torch.int64, device=device) + positions = torch.arange(past, past + num_tokens, dtype=torch.int64, device=device).view( + 1, -1 + ) + return self.model(ids, positions, md) + + def _after_sampling(self, state: RequestState): + """Read the sampler's choice and stopping decision off the request.""" + request = state._request + token = int(request.get_last_tokens(0)) + state.output_token_ids.append(token) + reason = None + if request.is_finished: + reason = request.py_finish_reason + if reason not in (FinishReason.END_ID, FinishReason.LENGTH): + raise RuntimeError(f"request {state.request_id} finished with reason {reason}") + state.finish_reason = reason + else: + state.pending_token = token + return token, reason + + def _record(self, state, step, past, num_tokens, source, token, reason, claimed=None) -> None: + cache = self.manager.kv_cache_map.get(state.request_id) + record = { + "step": step, + "past": past, + "tokens": num_tokens, + "positions": [past, past + num_tokens - 1], + "input": source, + "consumed": state.consumed, + "sampled": token, + "finish_reason": None if reason is None else reason.name, + "seq_slot": state._request.py_seq_slot, + } + if claimed is not None: + record["claimed"] = claimed + if cache is not None: + num_blocks = int(cache.num_blocks) + record.update( + capacity=int(cache.capacity), + history=int(cache.history_length), + committed=int(cache.num_committed_tokens), + pages=[ + np.frombuffer( + cache.get_base_page_indices(pool), dtype=np.int32, count=num_blocks + ).tolist() + for pool in range(len(self.manager.impl.layer_grouping)) + ], + ) + state.ledger.append(record) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/model.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/model.py new file mode 100644 index 000000000000..9b7ecaa2daaa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/model.py @@ -0,0 +1,683 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam six-block small model / sm_103 / tp1: token ids to logits, one request step per call. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no registered architecture: nothing in the engine builds this class, and it serves no +request on its own. A caller constructs it from a six-block configuration of the +checkpoint -- the checkpoint's own fields with ``num_hidden_layers = 6`` and the first six +entries of its two layer schedules -- loads its 2394 tensors and runs one step of one BF16 +request through it, with the attention batch state prepared on a six-layer +``KVCacheManagerV2``: ``prefill_context`` builds a fresh one for a single context, and the +stateful generator in ``generation.py`` keeps one across requests for cached decoding and +prefix reuse. It does not call the built-in Beam modeling. + +The model, all blocks with their own parameters and logical index:: + + x = embed_tokens[input_ids] [1, T, 5120] bf16, no scale or norm + x = block_0(x); x = block_1(x); x = block_2(x) dense, sliding attention + RoPE + x = block_3(x); x = block_4(x) sparse, sliding attention + RoPE + x = block_5(x) sparse, full attention, no RoPE + h = rmsnorm(x, norm) fp32 statistics and affine, one rounding + head = h @ lm_head^T [T, 151936] bf16, untied, no bias + logits = tanh(head.float() / 30) * 30 [1, T, 151936] fp32 + +The blocks are the dense and sparse block components, built for logical layers 0-2 and +3-5; each block's residual coefficients, attention kind and router are its own layer's. + +Steps. A call is one step of one request: a context step computes the T tokens of the +request's prompt that are not yet in its cache, after the ``past`` tokens that are (none +for a fresh context; a prefix reused from another request or computed by an earlier step +otherwise), and a generation step appends one token. The tokens sit at absolute positions +``past .. past + T - 1``; the returned rows are those tokens' logits. + +Cache. One call appends every layer's keys and values to one ``KVCacheManagerV2`` that +holds exactly logical layers 0-5, each in its own local slot: ``layer_offsets`` maps +logical layer i to slot i, and slot i's rows in the pool tables belong to it alone. The +five sliding layers attend through a 512-token window and layer 5 to the whole history; +the manager states this as its per-layer window vector ``[512] * 5 + [None]`` (``None`` is +the full ``max_seq_len``), and the attention op of each layer is called with that layer's +window. Before the first layer runs, ``check_cache`` verifies the manager, its slots and +the host pool tables the attention op reads -- that they address each slot's own key and +value pages as the manager allocated them -- and ``check_request`` verifies the request: +that the manager holds an active cache for it whose history ends exactly where the step +begins (a step never rewrites a cached token) and with the capacity the step needs, that +every page the step reads or writes is a valid page of its pool, and that the block-offset +table the op reads holds exactly the manager's addresses of those pages. Earlier tokens' keys +and values are read from the pool on a cached context step (paged-context FMHA) and on +every generation step. + +Host synchronization per call: one for the model's token-id range, position and +block-offset check, one for each block's position check, and each sparse block's MoE +dispatch (two, plus two per selected expert). A call is therefore not capturable in a CUDA +graph. +""" + +from __future__ import annotations + +import contextlib +import weakref +from collections import OrderedDict +from typing import Iterator, Mapping + +import numpy as np +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp +from tensorrt_llm.bindings import DataType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping as ParallelMapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX, LayerId, PageIndexMode + +from .. import _attention +from ..dense_block_bf16__sm_103__tp1 import block as dense_block +from ..moe_bf16__sm_103__tp1.moe import check_weights +from ..moe_bf16__sm_103__tp1.moe import weight_shapes as moe_weight_shapes +from ..sparse_block_bf16__sm_103__tp1 import block as sparse_block +from .ops.activation.beam_logit_softcap import beam_logit_softcap + +NUM_LAYERS = 6 +DENSE_LAYERS = (0, 1, 2) +SPARSE_LAYERS = (3, 4, 5) +#: The six-block schedules: the first six entries of the checkpoint's. +LAYER_TYPES = ("sliding_attention",) * 5 + ("full_attention",) +MLP_LAYER_TYPES = ("dense",) * 3 + ("sparse",) * 3 + +#: The configuration this component computes, beyond what its blocks check themselves. +#: Anything else is rejected at construction rather than computed approximately. +SUPPORTED_CONFIG = { + "num_hidden_layers": NUM_LAYERS, + "hidden_size": 5120, + "vocab_size": 151936, + "rms_norm_eps": 1e-6, + "tie_word_embeddings": False, + "final_logit_softcapping": 30.0, +} + +#: The KV cache ``prefill_context`` builds: bf16, 32-token pages, a 1024-token sequence +#: capacity and one request. A window equal to ``max_seq_len`` is the whole context. +KV_CACHE = {"tokens_per_block": 32, "max_seq_len": 1024, "max_tokens": 1024} +ATTENTION_WINDOWS = (512,) * 5 + (KV_CACHE["max_seq_len"],) + +#: The most tokens one call may compute: the sequence capacity. A request's whole KV +#: length is bounded by its cache's ``max_seq_len``. +MAX_TOKENS = KV_CACHE["max_seq_len"] + +_SM = (10, 3) + +#: Checkpoint names of the model's own tensors, beside its blocks'. +EMBED = "model.embed_tokens.weight" +NORM = "model.norm.weight" +HEAD = "lm_head.weight" + + +def block_prefix(layer_idx: int) -> str: + return f"model.layers.{layer_idx}." + + +def check_config(config) -> None: + """Raise ValueError naming the first field outside what this component computes.""" + for field, expected in SUPPORTED_CONFIG.items(): + actual = getattr(config, field, None) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this model computes {expected!r}") + for field, expected in (("layer_types", LAYER_TYPES), ("mlp_layer_types", MLP_LAYER_TYPES)): + actual = getattr(config, field, None) + if actual is None or tuple(actual) != expected: + raise ValueError( + f"unsupported {field}={actual!r}; this model computes {list(expected)}" + ) + for layer_idx in DENSE_LAYERS: + dense_block.check_config(config, layer_idx) + for layer_idx in SPARSE_LAYERS: + sparse_block.check_config(config, layer_idx) + + +def weight_schema(config) -> "OrderedDict[str, tuple[tuple[int, ...], torch.dtype]]": + """Every tensor by checkpoint name, in checkpoint order, with its shape and dtype. + + 2394 tensors: the embedding, each block's (15 per dense block, 782 per sparse block: + 12 attention and norm tensors and the MoE's FP32 router and 256 BF16 experts), the + final norm and the head. Matrices are [out, in]; the embedding is [vocab, hidden]. + """ + vocab, hidden = config.vocab_size, config.hidden_size + schema = OrderedDict([(EMBED, ((vocab, hidden), torch.bfloat16))]) + for layer_idx in range(NUM_LAYERS): + prefix = block_prefix(layer_idx) + if layer_idx in DENSE_LAYERS: + names = dense_block.weight_shapes(config).items() + else: + names = sparse_block.weight_shapes(config).items() + for name, shape in names: + schema[prefix + name] = (shape, torch.bfloat16) + if layer_idx in SPARSE_LAYERS: + for name, shape_dtype in moe_weight_shapes(config).items(): + schema[prefix + sparse_block.MLP_PREFIX + name] = shape_dtype + schema[NORM] = ((hidden,), torch.bfloat16) + schema[HEAD] = ((vocab, hidden), torch.bfloat16) + return schema + + +def check_cache(md: TrtllmAttentionMetadata, kv_len: int) -> None: + """Raise ValueError unless the call's cache holds the six layers in distinct valid slots. + + Required, before any layer writes: a ``KVCacheManagerV2`` holding exactly logical + layers 0-5, mapped one-to-one onto its six local slots; 8 bf16 kv heads of 128 in + every slot; a 512-token window on the sliding layers and the whole context on layer + 5, the full window holding the request's ``kv_len`` tokens; keys and values in one + shared page table per pool; host pool tables that address, for every slot, exactly + the key and value pages the manager allocated to that slot, so that no two slots + share storage (``_check_pool_tables``); and an int32 CUDA block-offset table of the + shape the attention op indexes. Reads host-side manager state and tensor metadata only. + """ + mgr = getattr(md, "kv_cache_manager", None) + if not isinstance(mgr, KVCacheManagerV2): + raise ValueError(f"the KV cache must be a KVCacheManagerV2, got {type(mgr).__name__}") + offsets = dict(getattr(mgr, "layer_offsets", None) or {}) + if set(offsets) != set(range(NUM_LAYERS)) or mgr.num_local_layers != NUM_LAYERS: + raise ValueError( + f"the KV cache must hold exactly logical layers 0-{NUM_LAYERS - 1}; it holds " + f"{sorted(offsets)} in {mgr.num_local_layers} local slots" + ) + slots = [offsets[layer_idx] for layer_idx in range(NUM_LAYERS)] + if any(type(slot) is not int for slot in slots) or sorted(slots) != list(range(NUM_LAYERS)): + raise ValueError(f"the six layers must map one-to-one onto slots 0-5, got {slots}") + heads = list(getattr(mgr, "num_kv_heads_per_layer", None) or []) + head_dims = list(getattr(mgr, "head_dim_per_layer", None) or []) + if ( + mgr.dtype != DataType.BF16 + or heads != [8] * NUM_LAYERS + or (head_dims and head_dims != [128] * NUM_LAYERS) + or mgr.head_dim != 128 + ): + raise ValueError( + f"every slot must hold 8 bf16 kv heads of 128; got {mgr.dtype}, {heads}, " + f"head_dim {mgr.head_dim}" + ) + windows = list(mgr.max_attention_window_vec) + for layer_idx, slot in enumerate(slots): + window = windows[slot] + if LAYER_TYPES[layer_idx] == "sliding_attention": + retained = mgr.max_seq_len if window is None else window + if retained != ATTENTION_WINDOWS[layer_idx]: + raise ValueError( + f"sliding layer {layer_idx} must keep a {ATTENTION_WINDOWS[layer_idx]}-token " + f"window in slot {slot}; the manager keeps {retained}" + ) + elif window is not None: + raise ValueError( + f"full-attention layer {layer_idx} must keep the whole context in slot {slot}; " + f"the manager keeps a {window}-token window" + ) + if mgr.max_seq_len < kv_len or md.max_seq_len != mgr.max_seq_len: + raise ValueError( + f"the cache holds {mgr.max_seq_len} tokens per sequence and the metadata states " + f"{md.max_seq_len}; the request reaches {kv_len}" + ) + _check_pool_tables(md, mgr, slots) + offsets = md.kv_cache_block_offsets + if ( + not isinstance(offsets, torch.Tensor) + or offsets.device.type != "cuda" + or offsets.dtype != torch.int32 + or offsets.dim() != 4 + or tuple(offsets.shape[i] for i in (0, 2, 3)) + != (mgr.num_attention_op_pools, 2, mgr.max_blocks_per_seq) + or offsets.shape[1] < 1 + ): + raise ValueError( + "the block-offset table must be a CUDA torch.int32 [pools, sequences, 2, " + f"{mgr.max_blocks_per_seq}] tensor with {mgr.num_attention_op_pools} pools" + ) + + +def _pool_windows(mgr: KVCacheManagerV2) -> dict: + """Each pool's attention window, from the windows of the slots it holds (``None``: full).""" + windows: dict = {} + for slot in range(NUM_LAYERS): + pool = int(mgr.impl.get_layer_group_id(LayerId(slot))) + window = mgr.max_attention_window_vec[slot] + if windows.setdefault(pool, window) != window: + raise ValueError(f"pool {pool} holds slots with windows {windows[pool]} and {window}") + return windows + + +def check_request(md: TrtllmAttentionMetadata, num_tokens: int, kv_len: int): + """Raise ValueError unless the call continues its request's cache on valid pages. + + The call reads keys and values of earlier tokens and writes its own ``num_tokens`` + after ``kv_len - num_tokens`` cached ones. Required, from the manager's own state: one + request id, whose cache in the manager is active, holds exactly those ``kv_len - + num_tokens`` tokens (its ``history_length``: the call appends right after them, so it + never rewrites a token the cache holds -- committed or reused prefix pages included, + which the manager keeps within the history -- and leaves no unwritten gap) and has + capacity for ``kv_len`` tokens; and for every pool, a valid page of that pool -- not + ``BAD_PAGE_INDEX``, inside the pool -- at every block ordinal a token of the call can + read or write: on the full layer's pool every ordinal below ``kv_len``, on the sliding + layers' pool those from the oldest key the call's first query attends to. Ordinals + older than that may carry the manager's out-of-window sentinel. + + Returns ``(expected, live)``: host int32 ``[pools, 2, n]`` block-offset entries the op + must read for those pages (the manager's page-to-offset conversion of each page) and + a bool ``[pools, n]`` mask of the ordinals checked, ``n = ceil(kv_len / tokens_per_block)``. + The caller compares them with the CUDA table. Reads host-side manager state only. + """ + mgr = md.kv_cache_manager + request_ids = list(md.request_ids or ()) + if len(request_ids) != 1: + raise ValueError(f"one request per call is supported; metadata names {request_ids}") + (request_id,) = request_ids + kv_cache = mgr.kv_cache_map.get(request_id) + if kv_cache is None or not kv_cache.is_active: + raise ValueError(f"request {request_id} holds no active cache in this manager") + history = int(kv_cache.history_length) + if history != kv_len - num_tokens: + raise ValueError( + f"request {request_id}'s cache holds {history} tokens; the call writes tokens " + f"{kv_len - num_tokens}..{kv_len - 1}, but a step must append right after the tokens " + "the cache holds" + ) + if kv_cache.capacity < kv_len: + raise ValueError( + f"request {request_id} has capacity for {kv_cache.capacity} tokens; the call reaches " + f"{kv_len}" + ) + tokens_per_block = mgr.tokens_per_block + num_blocks = -(-kv_len // tokens_per_block) + held = int(kv_cache.num_blocks) + if held < num_blocks: + raise ValueError(f"request {request_id} holds {held} blocks; the call needs {num_blocks}") + first_query = kv_len - num_tokens + pools = _pool_windows(mgr) + num_pools = len(mgr.impl.layer_grouping) + if sorted(pools) != list(range(num_pools)): + raise ValueError(f"the six slots occupy pools {sorted(pools)} of {num_pools}") + expected = torch.zeros((num_pools, 2, num_blocks), dtype=torch.int32) + live = torch.zeros((num_pools, num_blocks), dtype=torch.bool) + for pool, window in sorted(pools.items()): + oldest = 0 if window is None else max(0, first_query - window + 1) + pages = np.frombuffer(kv_cache.get_base_page_indices(pool), dtype=np.int32, count=held) + layer = next( + LayerId(slot) + for slot in range(NUM_LAYERS) + if int(mgr.impl.get_layer_group_id(LayerId(slot))) == pool + ) + scale, value_offset = int(mgr.index_scales[pool]), int(mgr.kv_offset[pool]) + bound = int(mgr.impl.get_page_index_upper_bound(layer, Role.KEY)) + for ordinal in range(oldest // tokens_per_block, num_blocks): + page = int(pages[ordinal]) + if page == BAD_PAGE_INDEX or page < 0 or scale * page + value_offset >= bound: + raise ValueError( + f"request {request_id} has page {page} at block {ordinal} of pool {pool}; " + f"the call reads or writes that block and the pool holds pages below " + f"{bound // scale}" + ) + expected[pool, 0, ordinal] = scale * page + expected[pool, 1, ordinal] = scale * page + value_offset + live[pool, ordinal] = True + return expected, live + + +def _overlap(spans) -> bool: + spans = sorted(spans) + return any(end > start for (_, end), (start, _) in zip(spans, spans[1:])) + + +def _check_pool_tables(md: TrtllmAttentionMetadata, mgr: KVCacheManagerV2, slots) -> None: + """Raise ValueError unless the host pool tables address each slot's own cache storage. + + The attention op reads slot s's row of ``host_kv_cache_pool_mapping`` as int32 + ``(pool, offset)`` and finds the slot's keys at ``host_kv_cache_pool_pointers[pool, 0] + + offset * 2 * page``, a page holding one slot's keys of ``tokens_per_block`` tokens and + the next page their values, with the pool's page indices from its block offsets. The + tables are host tensors a caller can replace, so each row is checked against the + manager's allocation itself: the pool is the slot's layer group, the offset keeps the + slot's key and value pages inside one slot of that pool, the address is where the + manager keeps the slot's keys, its values follow one page later, and no two slots' + pages, and no two pools, overlap. The five sliding layers and layer 5 are layer groups + of different slot sizes, which the manager keeps in separate pools; a cache whose layer + groups shared one pool is rejected rather than reasoned about. + """ + impl = mgr.impl + num_pools = len(impl.layer_grouping) + mapping = md.host_kv_cache_pool_mapping + pointers = md.host_kv_cache_pool_pointers + if ( + not isinstance(mapping, torch.Tensor) + or mapping.device.type != "cpu" + or mapping.dtype != torch.int32 + or tuple(mapping.shape) != (NUM_LAYERS, 2) + ): + raise ValueError(f"the pool mapping must be a host torch.int32 [{NUM_LAYERS}, 2] table") + if ( + not isinstance(pointers, torch.Tensor) + or pointers.device.type != "cpu" + or pointers.dtype != torch.int64 + or tuple(pointers.shape) != (num_pools, 2) + ): + raise ValueError( + f"the pool pointers must be a host torch.int64 [{num_pools}, 2] table, one row per pool" + ) + if ( + mgr.kv_cache_type != CacheTypeCpp.SELF + or mgr.enable_swa_scratch_reuse + or mgr.num_attention_op_pools != num_pools + ): + raise ValueError( + "the cache must hold keys and values with one shared page table per pool, " + "without sliding-window scratch reuse" + ) + page = mgr.tokens_per_block * 8 * 128 * torch.bfloat16.itemsize + pairs, pools = [], {} + for layer_idx, slot in enumerate(slots): + layer = LayerId(slot) + pool, offset = mapping[slot].tolist() + group = int(impl.get_layer_group_id(layer)) + scale = impl.get_page_index_scale(layer, Role.KEY) + if pool != group or not 0 <= offset < scale // 2: + raise ValueError( + f"layer {layer_idx} has pool row ({pool}, {offset}); the manager keeps slot {slot} " + f"in pool {group}, whose slots hold {scale // 2} key and value page pairs" + ) + start = impl.get_mem_pool_base_address(layer, Role.KEY, PageIndexMode.PER_LAYER) + key = impl.get_mem_pool_base_address(layer, Role.KEY, PageIndexMode.SHARED) + value = impl.get_mem_pool_base_address(layer, Role.VALUE, PageIndexMode.SHARED) + addressed = int(pointers[pool, 0]) + offset * 2 * page + if ( + addressed != key + or int(pointers[pool, 1]) != 0 + or value != key + page + or impl.get_page_stride(layer, Role.KEY) != page + or impl.get_page_stride(layer, Role.VALUE) != page + or not start <= key <= key + 2 * page <= start + scale * page + ): + raise ValueError( + f"layer {layer_idx}'s pool row ({pool}, {offset}) with pool pointers " + f"{pointers[pool].tolist()} addresses {addressed:#x}; the manager keeps slot " + f"{slot}'s keys at {key:#x} and values at {value:#x} in pool {pool} from " + f"{start:#x}, in {page}-byte pages" + ) + extent = (start, key + impl.get_page_index_upper_bound(layer, Role.KEY) * page) + if pools.setdefault(pool, extent) != extent: + raise ValueError( + f"pool {pool} spans {pools[pool]} for one slot and {extent} for another" + ) + pairs.append((key, key + 2 * page)) + if _overlap(pairs) or _overlap(pools.values()): + raise ValueError( + f"cache slots share storage: key and value pages {pairs}, pools {sorted(pools.values())}" + ) + + +@contextlib.contextmanager +def prefill_context(num_tokens: int, request_id: int = 0) -> Iterator[TrtllmAttentionMetadata]: + """A fresh six-layer ``KVCacheManagerV2`` and prepared metadata for one new request. + + The manager holds logical layers 0-5 in slots 0-5 (bf16, 8 kv heads of 128, + ``KV_CACHE`` pages and capacity, windows ``ATTENTION_WINDOWS``) with block reuse, + partial reuse, sliding-window scratch reuse and the overlap scheduler off; the + metadata holds one context request of ``num_tokens`` tokens starting at an empty + cache. The manager is shut down on exit, also after an error. + """ + mgr = KVCacheManagerV2( + KvCacheConfig( + max_tokens=KV_CACHE["max_tokens"], + max_attention_window=list(ATTENTION_WINDOWS), + enable_block_reuse=False, + enable_partial_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheTypeCpp.SELF, + num_layers=NUM_LAYERS, + num_kv_heads=8, + head_dim=128, + tokens_per_block=KV_CACHE["tokens_per_block"], + max_seq_len=KV_CACHE["max_seq_len"], + max_batch_size=1, + mapping=ParallelMapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.BF16, + disable_overlap_scheduler=True, + ) + try: + mgr.add_dummy_requests([request_id], token_nums=[num_tokens]) + md = TrtllmAttentionMetadata( + max_num_requests=1, + max_num_tokens=MAX_TOKENS, + kv_cache_manager=mgr, + num_contexts=1, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0]), + seq_lens=torch.tensor([num_tokens], dtype=torch.int), + request_ids=[request_id], + prompt_lens=[num_tokens], + ) + md.prepare() + yield md + finally: + mgr.shutdown() + + +class BeamSmallModelPrefillBf16Sm103Tp1(nn.Module): + """Beam six-block model: token ids in, every token's FP32 vocabulary logits out. + + State: ``embed_tokens`` ``[151936, 5120]`` and ``lm_head`` ``[151936, 5120]``, two + separate BF16 tensors; ``layers``, dense blocks 0-2 and sparse blocks 3-5, each with + its own tensors, router and 256 experts; ``norm`` ``[5120]`` BF16. ``load_weights`` + loads all 2394 tensors by checkpoint name; the model runs only after that succeeded. + + ``forward(input_ids [1, T] int32/int64, position_ids [1, T], attn_metadata)`` returns + ``[1, T, 151936]`` ``torch.float32`` on the model's device: row t, entry v is the + softcapped logit for vocabulary id v of the t-th token of the call. Accepted: + 1 <= T <= ``MAX_TOKENS`` token ids in ``[0, 151936)``, contiguous on the model's device; + the prepared ``TrtllmAttentionMetadata`` of one step of one request -- a context step + of T tokens after ``past`` cached ones (paged-context FMHA when ``past > 0``) or a + one-token generation step -- on the six-layer cache ``check_cache`` describes, its + request's pages valid as ``check_request`` describes and its KV length at most the + cache's ``max_seq_len``; and integer positions equal to ``arange(past, past + T)``. + ``prefill_context(T)`` yields such metadata for a fresh context; the generator in + ``generation.py`` for every step of its requests. Anything else is rejected with + ValueError before the embedding lookup or any other kernel runs, including an id + outside the vocabulary. Inputs and parameters are never modified. Batches of several + requests and padding are not computed here. + + Weights cannot change under a cache built from them: ``load_weights`` raises while a + generator holding this model's cache is open (``close`` it first). + """ + + def __init__(self, config, device: str | torch.device = "cuda"): + super().__init__() + check_config(config) + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this model is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.eps = config.rms_norm_eps + self.logit_softcap = config.final_logit_softcapping + self._schema = weight_schema(config) + + def empty(name: str) -> nn.Parameter: + shape, dtype = self._schema[name] + return nn.Parameter(torch.empty(shape, dtype=dtype, device=device), requires_grad=False) + + self.embed_tokens = empty(EMBED) + self.layers = nn.ModuleList( + dense_block.BeamDenseBlockBf16Sm103Tp1(config, layer_idx, device=device) + if layer_idx in DENSE_LAYERS + else sparse_block.BeamSparseBlockBf16Sm103Tp1(config, layer_idx, device=device) + for layer_idx in range(NUM_LAYERS) + ) + self.norm = empty(NORM) + self.lm_head = empty(HEAD) + self._loaded = False + #: Generators whose KV cache holds keys and values computed with these weights. + self._cache_holders: "weakref.WeakSet" = weakref.WeakSet() + + @property + def device(self) -> torch.device: + return self.embed_tokens.device + + def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": + """All 2394 tensors by checkpoint name, in checkpoint order.""" + tensors = OrderedDict([(EMBED, self.embed_tokens)]) + for layer_idx, block in enumerate(self.layers): + for name, tensor in block.checkpoint_tensors().items(): + tensors[block_prefix(layer_idx) + name] = tensor + tensors[NORM] = self.norm + tensors[HEAD] = self.lm_head + return tensors + + def check_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Raise unless ``weights`` holds exactly the 2394 names with their dtypes and shapes. + + Reads tensor metadata only, never values, and copies nothing. + """ + check_weights(self._schema, weights, "model") + + @torch.no_grad() + def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Copy exactly the model's 2394 tensors, by checkpoint name, [out, in]. + + Names are ``model.embed_tokens.weight``, ``model.layers..``, + ``model.norm.weight`` and ``lm_head.weight``; dtypes are BF16 except each sparse + block's FP32 ``mlp.gate.weight`` and ``mlp.gate.e_score_correction_bias``; host or + CUDA tensors are copied as given, without transposing or converting. Every tensor + is checked before any parameter changes, so a rejected load leaves the model as it + was. The model runs only after a load completed: a failure while copying leaves it + unloaded until a later load succeeds. Copies proceed one tensor (one expert, one + block's attention tensors) at a time; no copy of a whole bank is made. + + Raises RuntimeError, before checking or copying anything, while an open generator + holds a KV cache computed with the current weights: its cached and reusable keys + and values would no longer match them. + """ + holders = [holder for holder in self._cache_holders if not holder.closed] + if holders: + raise RuntimeError( + f"{len(holders)} open generator(s) hold KV caches computed with these weights; " + "close them before loading new weights" + ) + self.check_weights(weights) + self._loaded = False + for name, param in ((EMBED, self.embed_tokens), (NORM, self.norm), (HEAD, self.lm_head)): + param.copy_(weights[name]) + for layer_idx, block in enumerate(self.layers): + prefix = block_prefix(layer_idx) + block.load_weights( + {name[len(prefix) :]: w for name, w in weights.items() if name.startswith(prefix)} + ) + self._loaded = True + + def register_cache_holder(self, holder) -> None: + """Record a generator whose KV cache holds keys and values computed with these weights. + + ``holder`` must expose a ``closed`` attribute; while it is False ``load_weights`` + raises. Holders are referenced weakly. + """ + self._cache_holders.add(holder) + + def _check_call( + self, input_ids: torch.Tensor, position_ids: torch.Tensor, md: TrtllmAttentionMetadata + ) -> int: + """Raise ValueError on anything the model would otherwise compute wrongly; return T.""" + if not self._loaded: + raise ValueError("load_weights must complete before the model runs") + if ( + not isinstance(input_ids, torch.Tensor) + or input_ids.dtype not in (torch.int32, torch.int64) + or input_ids.layout != torch.strided + ): + raise ValueError("input_ids must be a dense torch.int32 or torch.int64 tensor") + if input_ids.dim() != 2 or input_ids.shape[0] != 1: + raise ValueError(f"input_ids must be [1, T], got {tuple(input_ids.shape)}") + num_tokens = input_ids.shape[1] + if not 1 <= num_tokens <= MAX_TOKENS: + raise ValueError( + f"supported token counts per call are 1..{MAX_TOKENS}, got {num_tokens}" + ) + if input_ids.device != self.device or not input_ids.is_contiguous(): + raise ValueError(f"input_ids must be contiguous on the model's device {self.device}") + if ( + not isinstance(position_ids, torch.Tensor) + or position_ids.dtype not in (torch.int32, torch.int64) + or tuple(position_ids.shape) != (1, num_tokens) + or position_ids.device != self.device + or not position_ids.is_contiguous() + ): + raise ValueError( + f"position_ids must be a contiguous integer [1, {num_tokens}] tensor on {self.device}" + ) + if not isinstance(md, TrtllmAttentionMetadata): + raise ValueError( + f"attn_metadata must be TrtllmAttentionMetadata, got {type(md).__name__}" + ) + mgr = getattr(md, "kv_cache_manager", None) + if not isinstance(mgr, KVCacheManagerV2): + raise ValueError(f"the KV cache must be a KVCacheManagerV2, got {type(mgr).__name__}") + pasts = { + _attention.check_metadata( + md, num_tokens, layer_idx=layer_idx, num_kv_heads=8, head_dim=128 + ) + for layer_idx in range(NUM_LAYERS) + } + (past,) = pasts + kv_len = past + num_tokens + check_cache(md, kv_len) + expected_table, live = check_request(md, num_tokens, kv_len) + table = md.kv_cache_block_offsets[:, 0, :, : expected_table.shape[-1]] + table_ok = ( + (table == expected_table.to(table.device, non_blocking=True)) + | ~live.to(table.device, non_blocking=True)[:, None, :] + ).all() + expected = torch.arange( + past, past + num_tokens, dtype=position_ids.dtype, device=self.device + ) + low, high = torch.aminmax(input_ids) + positions_ok, ids_ok, table_ok = torch.stack( + ( + (position_ids[0] == expected).all(), + (low >= 0) & (high < self.vocab_size), + table_ok, + ) + ).tolist() + if not positions_ok: + raise ValueError( + f"the call appends tokens {past}..{kv_len - 1} after {past} cached tokens; its " + "positions must be exactly those" + ) + if not table_ok: + raise ValueError( + "the block-offset table does not address the pages the manager allocated to " + "the request" + ) + if not ids_ok: + raise ValueError( + f"token ids must be in [0, {self.vocab_size}); got values in " + f"[{int(low)}, {int(high)}]" + ) + return num_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attn_metadata: TrtllmAttentionMetadata, + ) -> torch.Tensor: + num_tokens = self._check_call(input_ids, position_ids, attn_metadata) + x = nn.functional.embedding(input_ids, self.embed_tokens) + for block in self.layers: + x = block(x, position_ids, attn_metadata) + h = flashinfer_rmsnorm(x.view(num_tokens, self.hidden_size), self.norm, self.eps) + head = cublas_mm(h, self.lm_head.t()) + logits = beam_logit_softcap(head, self.logit_softcap) + return logits.view(1, num_tokens, self.vocab_size) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/__init__.py new file mode 100644 index 000000000000..23e614d4422b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this component. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/activation/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/activation/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/activation/beam_logit_softcap.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/activation/beam_logit_softcap.py new file mode 100644 index 000000000000..af45a7e6e391 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/activation/beam_logit_softcap.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Final-logit tanh softcap in fp32: `tanh(logits.float() / cap) * cap`.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The head's output: every token's 151936 vocabulary logits, bf16, cap 30. +_HEAD = dict(vocab=151936, cap=30.0, dtype=torch.bfloat16) + + +class BeamLogitSoftcap(OpWrapper): + """Softcap the head projection `[T, V]`: `(logits.float() / cap).tanh() * cap`, fp32. + + The head's activation-dtype output is widened to fp32, divided by `cap`, passed + through tanh and multiplied by `cap`, each step in fp32; the fp32 result is + returned, in `(-cap, cap)` for finite inputs and exactly `+-cap` once tanh + saturates. Applied once, after the head projection and nothing else. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="head, 1 token", spec=dict(_HEAD, tokens=1)), + Cell(why="head, 17 tokens", spec=dict(_HEAD, tokens=17)), + Cell(why="head, 129 tokens: the longest context", spec=dict(_HEAD, tokens=129)), + Cell( + why="signed values from 0 to beyond saturation, -0 included", + spec=dict(_HEAD, tokens=1, saturating=True), + ), + ) + + note = """ + Not idempotent: a second application moves every nonzero logit toward zero + (`30 * tanh(30 * tanh(x / 30) / 30)` < `30 * tanh(x / 30)` for x > 0), and + omitting it changes every logit whose magnitude exceeds about 0.05. Both are + therefore visible on ordinary logits by comparing the returned values with this + expression of the captured head output, without any saturating input. + """ + + def __call__(self, logits: torch.Tensor, cap: float) -> torch.Tensor: + return (logits.float() / cap).tanh() * cap + + def reference(self, logits: torch.Tensor, cap: float) -> torch.Tensor: + """The same three fp32 steps, written out one by one.""" + widened = logits.to(torch.float32) + scaled = torch.div(widened, cap) + return torch.mul(torch.tanh(scaled), cap) + + def is_valid(self, logits: torch.Tensor, cap: float) -> None: + # A tensor cap broadcasts and a non-positive cap flips or removes the cap; an + # fp32 input that was rounded upstream is indistinguishable here. None raises. + assert isinstance(cap, float) and cap > 0, f"cap must be a positive float, got {cap!r}" + assert logits.dim() == 2, "logits must be [T, vocab]" + assert logits.dtype == torch.bfloat16, "logits must be the bf16 head output" + + +beam_logit_softcap = BeamLogitSoftcap() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/__init__.py new file mode 100644 index 000000000000..c10dd6cb1a2e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam decoder block 3 (sparse MoE, sliding attention) / sm_103 / tp1.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/block.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/block.py new file mode 100644 index 000000000000..7746f948c39b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/block.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam decoder blocks 3-5 / sm_103 / tp1: one sparse (MoE) block, constructed directly. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine builds this class, and it serves no request. +A caller constructs it from the checkpoint's configuration for one of the sparse blocks +3, 4 or 5, loads the block's 782 tensors and runs one step of one BF16 request through it +-- a context step after any cached prefix, or a one-token generation step -- with the +attention batch state prepared on a ``KVCacheManagerV2`` that holds the block's logical +layer. The forward is assembled from catalog entries, the attention sublayer +shared with the dense block component (``../_attention.py``) and the MoE component +(``../moe_bf16__sm_103__tp1``); it does not call the built-in Beam modeling and owns no +dense MLP. + +The block, with block 3's residual coefficients (block i's are in ``depth_scales``):: + + h = rmsnorm(x, input_layernorm) + a = attention(h) see ../_attention.py + x = rmsnorm(a, post_attention_layernorm) / sqrt(8) + x * sqrt(7) / sqrt(8) + h = rmsnorm(x, pre_feedforward_layernorm) + m = moe(h) routed scaling 2.5 included + out = rmsnorm(m, post_feedforward_layernorm) / sqrt(9) + x * sqrt(8) / sqrt(9) + +Blocks 3 and 4 attend through the 512-token sliding window with NeoX RoPE; block 5 is a +full-attention layer without position embedding: no rotation of q or k, the per-token +query scale on q, and the whole context as its window. + +Every RMSNorm computes its statistics, normalization and affine in fp32 and rounds once; +each residual add rounds both scaled terms to bf16, then their sum. The residual +coefficients are the logical block's: block i's attention add is the 2(i + 1)-th add of +the residual stream, its feed-forward add the next one. + +Logical layer and cache slot: ``layer_idx`` fixes the block's semantics; the cache slot is +the call's manager's local index for that layer -- 0 for a manager whose ``layer_mask`` +selects only it, its own position in a manager holding several layers. A manager that +does not hold the layer -- a default one-layer manager holds layer 0 -- is rejected +before any kernel runs. + +Host synchronization per call: one for the position check, and the MoE's dispatch +synchronizations (two, plus two per selected expert). Positions are absolute: a sliding +layer attends to the newest 512 keys of the request's history up to each query, the +full-attention layer to all of them. +""" + +from __future__ import annotations + +import math +from collections import OrderedDict +from typing import Mapping + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.residual.beam_depth_residual import ( + beam_depth_residual, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + +from .. import _attention +from ..moe_bf16__sm_103__tp1.moe import BeamMoeBf16Sm103Tp1, check_weights +from ..moe_bf16__sm_103__tp1.moe import check_config as check_moe_config + +#: The configuration this component computes, as the checkpoint's config.json states it. +#: Anything else is rejected at construction rather than computed approximately; the +#: MoE component checks its own fields. +SUPPORTED_CONFIG = { + "hidden_size": 5120, + "num_attention_heads": 48, + "num_key_value_heads": 8, + "head_dim": 128, + "rms_norm_eps": 1e-6, + "attention_bias": False, + "attention_dropout": 0.0, + "sliding_window": 512, +} +SUPPORTED_LAYER_IDX = 3 +#: The sparse blocks of the checkpoint's schedule and the attention each one runs. +ATTENTION_TYPES = {3: "sliding_attention", 4: "sliding_attention", 5: "full_attention"} +SUPPORTED_LAYER_INDICES = tuple(ATTENTION_TYPES) +#: RoPE of the sliding layers; the full-attention layers have none. +SUPPORTED_ROPE = {"rope_type": "default", "rope_theta": 10000.0} + +#: The most tokens one call may append: the six-layer model's sequence capacity. Longer +#: calls are rejected; the request's whole KV length is bounded by its cache's capacity. +MAX_TOKENS = 1024 + +_SM = (10, 3) + +#: The MoE's tensors carry this prefix in the block's checkpoint names. +MLP_PREFIX = "mlp." + + +def weight_shapes(config) -> "OrderedDict[str, tuple[int, ...]]": + """The block's twelve attention and norm tensors by checkpoint name, [out, in], all BF16.""" + hidden = config.hidden_size + q_width = config.num_attention_heads * config.head_dim + kv_width = config.num_key_value_heads * config.head_dim + return OrderedDict( + [ + ("input_layernorm.weight", (hidden,)), + ("self_attn.gate_bias", (config.num_attention_heads,)), + ("self_attn.q_proj.weight", (q_width, hidden)), + ("self_attn.k_proj.weight", (kv_width, hidden)), + ("self_attn.v_proj.weight", (kv_width, hidden)), + ("self_attn.gate_proj.weight", (q_width, hidden)), + ("self_attn.o_proj.weight", (hidden, q_width)), + ("self_attn.q_norm.weight", (config.head_dim,)), + ("self_attn.k_norm.weight", (config.head_dim,)), + ("pre_feedforward_layernorm.weight", (hidden,)), + ("post_attention_layernorm.weight", (hidden,)), + ("post_feedforward_layernorm.weight", (hidden,)), + ] + ) + + +def _attribute(name: str) -> str: + return name.removesuffix(".weight").replace(".", "_") + + +def depth_scales(layer_idx: int, is_attention: bool) -> tuple[float, float]: + """(branch, residual) scales of a block's attention or feed-forward residual add. + + The n-th add of the residual stream scales the branch by 1 / sqrt(n) and the stream + by sqrt(n - 1) / sqrt(n); block i's attention add is the 2(i + 1)-th, its + feed-forward add the next one. + """ + num_adds = 2 * (layer_idx + 1) + (0 if is_attention else 1) + denom = math.sqrt(num_adds) + return 1.0 / denom, math.sqrt(num_adds - 1) / denom + + +def check_config(config, layer_idx: int) -> None: + """Raise ValueError naming the first field outside what this component computes.""" + if type(layer_idx) is not int or layer_idx not in SUPPORTED_LAYER_INDICES: + raise ValueError( + f"only sparse blocks {SUPPORTED_LAYER_INDICES} are supported, got layer_idx={layer_idx!r}" + ) + for field, expected in SUPPORTED_CONFIG.items(): + actual = getattr(config, field, None) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this block computes {expected!r}") + attention_type = ATTENTION_TYPES[layer_idx] + layer_types = getattr(config, "layer_types", None) or [] + if len(layer_types) <= layer_idx or layer_types[layer_idx] != attention_type: + raise ValueError( + f"block {layer_idx} must be a {attention_type} layer: " + f"{layer_types[layer_idx] if len(layer_types) > layer_idx else None!r}" + ) + rope_parameters = getattr(config, "rope_parameters", None) or {} + rope = rope_parameters.get(attention_type) + expected_rope = SUPPORTED_ROPE if attention_type == "sliding_attention" else None + if rope != expected_rope or attention_type not in rope_parameters: + raise ValueError( + f"unsupported {attention_type} rope parameters {rope!r}; need {expected_rope!r}" + ) + check_moe_config(config, layer_idx) + + +class BeamSparseBlockBf16Sm103Tp1(nn.Module): + """Beam sparse decoder block 3, 4 or 5: BF16 weights and activations, an FP32 router, one GPU, + one request step. + + State: the twelve BF16 attention and norm tensors of ``weight_shapes`` and ``mlp``, the + MoE component with its router and 256 experts; ``load_weights`` loads all 782. + + ``forward(hidden_states [1, T, 5120] bf16, position_ids [1, T], attn_metadata)`` + returns ``[1, T, 5120]`` bf16 for 1 <= T <= ``MAX_TOKENS``, a prepared + ``TrtllmAttentionMetadata`` holding one step of one request on a BF16 + ``KVCacheManagerV2`` (8 kv heads of 128) that holds the block's logical layer -- a + context step of T tokens after ``past`` cached ones (paged-context FMHA when + ``past > 0``) or a one-token generation step -- and positions ``arange(past, past + T)``. + Inputs outside that are rejected before any kernel runs; batches of several requests + and other blocks are not computed here. Inputs and parameters are never modified. + """ + + def __init__( + self, config, layer_idx: int = SUPPORTED_LAYER_IDX, device: str | torch.device = "cuda" + ): + super().__init__() + check_config(config, layer_idx) + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this block is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.layer_idx = layer_idx + self.eps = config.rms_norm_eps + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.is_sliding = ATTENTION_TYPES[layer_idx] == "sliding_attention" + self.window = config.sliding_window if self.is_sliding else None + self.rope_theta = SUPPORTED_ROPE["rope_theta"] if self.is_sliding else None + self.attention_scales = depth_scales(layer_idx, is_attention=True) + self.feed_forward_scales = depth_scales(layer_idx, is_attention=False) + self._shapes = weight_shapes(config) + for name, shape in self._shapes.items(): + tensor = torch.empty(shape, dtype=torch.bfloat16, device=device) + self.register_parameter(_attribute(name), nn.Parameter(tensor, requires_grad=False)) + self.mlp = BeamMoeBf16Sm103Tp1(config, layer_idx, device=device) + self._loaded = False + + @property + def hidden_size(self) -> int: + return self._shapes["input_layernorm.weight"][0] + + def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": + """The block's 782 tensors by checkpoint name: its own twelve, then the MoE's.""" + tensors = OrderedDict((name, getattr(self, _attribute(name))) for name in self._shapes) + for name, tensor in self.mlp.checkpoint_tensors().items(): + tensors[MLP_PREFIX + name] = tensor + return tensors + + @torch.no_grad() + def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Copy exactly the block's 782 tensors, by checkpoint name, [out, in]. + + The twelve attention and norm tensors are BF16; the ``mlp.``-prefixed MoE tensors + are as ``BeamMoeBf16Sm103Tp1.load_weights`` describes (FP32 router, BF16 experts). + Every tensor is checked before any parameter changes, so a rejected load leaves + the block as it was. The block runs only after a load completed: a failure while + copying leaves it unloaded. + """ + if not isinstance(weights, Mapping): + raise TypeError(f"weights must be a mapping, got {type(weights).__name__}") + own = {n: w for n, w in weights.items() if not str(n).startswith(MLP_PREFIX)} + moe = {n[len(MLP_PREFIX) :]: w for n, w in weights.items() if str(n).startswith(MLP_PREFIX)} + check_weights( + OrderedDict((n, (shape, torch.bfloat16)) for n, shape in self._shapes.items()), + own, + "block", + ) + self.mlp.check_weights(moe) + self._loaded = False + device = self.input_layernorm.device + staged = {name: own[name].to(device=device, copy=True) for name in self._shapes} + for name, tensor in staged.items(): + getattr(self, _attribute(name)).copy_(tensor) + self.mlp.load_weights(moe) + self._loaded = True + + def _check_step( + self, hidden_states: torch.Tensor, position_ids: torch.Tensor, md: TrtllmAttentionMetadata + ) -> None: + """Raise ValueError on anything this block would otherwise compute wrongly.""" + if not self._loaded: + raise ValueError("load_weights must be called before the block runs") + _attention.check_step( + hidden_states, + position_ids, + md, + layer_idx=self.layer_idx, + hidden_size=self.hidden_size, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + device=self.input_layernorm.device, + max_tokens=MAX_TOKENS, + ) + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """Project the prepared metadata onto thop_attention's explicit batch state.""" + return _attention.step_args(md, self.layer_idx) + + def attention(self, h: torch.Tensor, positions: torch.Tensor, step: dict) -> torch.Tensor: + """Normalized input [T, hidden] -> attention output [T, hidden], before post_attention_layernorm.""" + return _attention.attention( + h, + positions, + step, + q_proj=self.self_attn_q_proj, + k_proj=self.self_attn_k_proj, + v_proj=self.self_attn_v_proj, + gate_proj=self.self_attn_gate_proj, + o_proj=self.self_attn_o_proj, + q_norm=self.self_attn_q_norm, + k_norm=self.self_attn_k_norm, + gate_bias=self.self_attn_gate_bias, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + eps=self.eps, + rope_theta=self.rope_theta, + window=self.window if self.is_sliding else step["max_seq_len"], + ) + + @torch.no_grad() + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + attn_metadata: TrtllmAttentionMetadata, + ) -> torch.Tensor: + self._check_step(hidden_states, position_ids, attn_metadata) + step = self.step_args(attn_metadata) + num_tokens = hidden_states.shape[1] + x = hidden_states.view(num_tokens, self.hidden_size) + positions = position_ids.view(num_tokens) + + h = flashinfer_rmsnorm(x, self.input_layernorm, self.eps) + h = self.attention(h, positions, step) + h = flashinfer_rmsnorm(h, self.post_attention_layernorm, self.eps) + x = beam_depth_residual(h, x, *self.attention_scales) + + h = flashinfer_rmsnorm(x, self.pre_feedforward_layernorm, self.eps) + h = self.mlp(h) + h = flashinfer_rmsnorm(h, self.post_feedforward_layernorm, self.eps) + x = beam_depth_residual(h, x, *self.feed_forward_scales) + return x.view(1, num_tokens, self.hidden_size) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index e465eda7d093..915917086cf1 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -320,6 +320,40 @@ def _get_workspace_size( return max(context_size, generation_size) +@lru_cache(maxsize=64) +def _fallback_has_paged_context_kernel( + head_size: int, kv_cache_dtype: DataType, tokens_per_block: int +) -> bool: + kernel_exists = getattr(thop, "fused_context_fmha_kernel_exists", None) + if kernel_exists is None: + return True + return bool( + kernel_exists( + head_size=head_size, + kv_cache_dtype=kv_cache_dtype, + tokens_per_block=tokens_per_block, + output_dtype=kv_cache_dtype, + ) + ) + + +def _fallback_serves_paged_context( + head_size: int, meta: "TrtllmAttentionMetadata", kv_cache_dtype: Optional[DataType] +) -> bool: + """Whether the fallback (thop) FMHA can run this 16-bit context phase. + + With ``use_paged_context_fmha`` (chunked prefill, KV block reuse or draft tokens) + the context attends to cached KV, which needs a fused paged-context kernel. A build + can lack it for a head size and page size, and the thop op then refuses the request, + so such a context is not a candidate for the fallback. Builds whose bindings predate + the ``fused_context_fmha_kernel_exists`` lookup keep the previous selection. + """ + if not meta.use_paged_context_fmha or meta.tokens_per_block is None: + return True + dtype = kv_cache_dtype if kv_cache_dtype is not None else DataType.BF16 + return _fallback_has_paged_context_kernel(head_size, dtype, meta.tokens_per_block) + + class FlashInferTrtllmGenFmha(PhasedFmha): """ An attention backend using pure trtllm-gen kernels from flashinfer. @@ -636,6 +670,7 @@ def _is_supported_with_reason( meta.use_paged_context_fmha and meta.tokens_per_block not in self.SUPPORTED_TOKENS_PER_BLOCK ) + and _fallback_serves_paged_context(attn.head_dim, meta, kv_cache_dtype) ): # NVBug 6579626: the per-layer host overhead of the FlashInfer # TRTLLM-Gen context path regresses TTFT for small BF16 batches. diff --git a/tensorrt_llm/_torch/models/modeling_beam.py b/tensorrt_llm/_torch/models/modeling_beam.py new file mode 100644 index 000000000000..1b41258d71d8 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_beam.py @@ -0,0 +1,853 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Beam decoder components for the PyTorch backend. + +Beam attention is grouped-query attention with an explicit ``head_dim`` (the query +width is independent of ``hidden_size``), learned per-head RMSNorm on Q and K, NeoX +RoPE on sliding-window layers, Scalable-Softmax query scaling on full (NoPE) layers, +and a sigmoid output gate with a per-head learned bias. Decoder blocks use sandwich +norms (four RMSNorms) and depth-scaled residual adds. Sparse layers route with a +sigmoid grouped top-k router whose fp32 correction bias only affects selection. + +The transforms reproduce the reference model's dtype staging exactly: + +* RMSNorm computes the statistics, normalization and affine multiply in fp32 and + rounds once to the activation dtype. +* RoPE builds cos/sin in fp32 and rotates in fp32 before one cast; SSMax scales the + query in fp32 before one cast. +* The output gate adds its bias, applies the sigmoid and multiplies in the + activation dtype. +* SwiGLU applies SiLU(gate) * up in fp32 and casts before the down projection. +* The router GEMM, scores, selection and weights stay fp32. Each weighted expert + route is rounded to the activation dtype, routes are accumulated in fp32 in + ascending expert order, and the routed scaling factor is applied in fp32. +* The depth residual multiplies and adds in the activation dtype. +* The logit softcap runs in fp32 on the LM-head logits. + +Projections, the paged KV cache and the attention core run through the standard +``Attention`` module and the selected attention backend; Q/K/V and the gate are +packed into the fused projection (per head ``[q_h, gate_h]``, then K, then V). The +cache holds post-norm, post-RoPE keys and projected values. +""" + +import math +from typing import Dict, List, NamedTuple, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig + +from ..attention.attention import Attention +from ..attention.backends import AttentionMetadata +from ..model_config import ModelConfig +from ..modules.decoder_layer import DecoderLayer +from ..modules.embedding import Embedding +from ..modules.linear import Linear, TensorParallelMode +from .modeling_utils import DecoderModel, DecoderModelForCausalLM + +BEAM_LAYER_TYPES = ("sliding_attention", "full_attention") +BEAM_MLP_LAYER_TYPES = ("dense", "sparse") +BEAM_ROPE_TYPES = ("default", "linear") +# Router parameters that keep their fp32 values regardless of the activation dtype. +BEAM_FP32_PARAMETER_SUFFIXES = ("mlp.gate.weight", "mlp.gate.e_score_correction_bias") + + +def _default_layer_types(num_hidden_layers: int) -> List[str]: + layer_types = [ + "sliding_attention" if (i + 1) % 6 else "full_attention" for i in range(num_hidden_layers) + ] + if layer_types: + layer_types[-1] = "full_attention" + return layer_types + + +def _default_mlp_layer_types(num_hidden_layers: int) -> List[str]: + num_dense = min(3, num_hidden_layers) + return ["dense"] * num_dense + ["sparse"] * (num_hidden_layers - num_dense) + + +def _normalize_rope_parameters(raw: Optional[dict]) -> dict: + """Validate Beam's sliding-RoPE / full-NoPE rope parameters.""" + if raw is None: + return { + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + } + if not isinstance(raw, dict) or set(raw) != {"sliding_attention", "full_attention"}: + raise ValueError( + "rope_parameters must contain exactly 'sliding_attention' and 'full_attention': " + f"{raw!r}" + ) + if raw["full_attention"] is not None: + raise ValueError("Beam full-attention layers run without rotary embeddings") + sliding = raw["sliding_attention"] + if not isinstance(sliding, dict): + raise ValueError("sliding_attention rope parameters must be a dictionary") + rope_type = sliding.get("rope_type", "default") + if rope_type not in BEAM_ROPE_TYPES: + raise ValueError(f"rope_type must be one of {BEAM_ROPE_TYPES}: {rope_type!r}") + theta = sliding.get("rope_theta", 10000.0) + if isinstance(theta, bool) or not isinstance(theta, (int, float)): + raise ValueError(f"rope_theta must be a number: {theta!r}") + normalized = {"rope_type": rope_type, "rope_theta": float(theta)} + if rope_type == "linear": + factor = sliding.get("factor") + if isinstance(factor, bool) or not isinstance(factor, (int, float)) or factor <= 0: + raise ValueError(f"linear RoPE needs a positive factor: {factor!r}") + normalized["factor"] = float(factor) + return {"sliding_attention": normalized, "full_attention": None} + + +class BeamConfig(PretrainedConfig): + """Architecture configuration of a Beam checkpoint (``config.json`` fields).""" + + model_type = "beam" + + def __init__( + self, + vocab_size: int = 151936, + hidden_size: int = 5120, + intermediate_size: int = 20480, + num_hidden_layers: int = 52, + num_attention_heads: int = 48, + num_key_value_heads: int = 8, + head_dim: int = 128, + hidden_act: str = "silu", + max_position_embeddings: int = 1048576, + rms_norm_eps: float = 1e-6, + rope_parameters: Optional[dict] = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + sliding_window: Optional[int] = 512, + moe_intermediate_size: int = 2560, + num_experts_per_tok: int = 8, + num_experts: int = 256, + norm_topk_prob: bool = True, + n_group: int = 8, + topk_group: int = 4, + routed_scaling_factor: float = 2.5, + final_logit_softcapping: Optional[float] = 30.0, + layer_types: Optional[List[str]] = None, + mlp_layer_types: Optional[List[str]] = None, + tie_word_embeddings: bool = False, + **kwargs, + ): + dtype = kwargs.pop("dtype", None) + torch_dtype = kwargs.pop("torch_dtype", None) + dtype = dtype if dtype is not None else torch_dtype + if dtype is None: + dtype = torch.bfloat16 + if isinstance(dtype, str): + dtype = getattr(torch, dtype) + kwargs.setdefault("architectures", ["BeamForCausalLM"]) + super().__init__(tie_word_embeddings=tie_word_embeddings, dtype=dtype, **kwargs) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + self.rope_parameters = _normalize_rope_parameters(rope_parameters) + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.sliding_window = sliding_window + self.moe_intermediate_size = moe_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.norm_topk_prob = norm_topk_prob + self.n_group = n_group + self.topk_group = topk_group + self.routed_scaling_factor = routed_scaling_factor + self.final_logit_softcapping = final_logit_softcapping + self.tie_word_embeddings = tie_word_embeddings + self.layer_types = ( + list(layer_types) if layer_types else _default_layer_types(num_hidden_layers) + ) + self.mlp_layer_types = ( + list(mlp_layer_types) + if mlp_layer_types + else _default_mlp_layer_types(num_hidden_layers) + ) + if getattr(self, "torch_dtype", None) != dtype: + try: + self.torch_dtype = dtype + except AttributeError: + # A read-only torch_dtype property derives from ``dtype``. + self.dtype = dtype + if self.torch_dtype != dtype: + raise ValueError(f"BeamConfig dtype {self.torch_dtype!r} != requested {dtype!r}") + self._validate() + + def _validate(self) -> None: + if len(self.layer_types) != self.num_hidden_layers: + raise ValueError("layer_types length must equal num_hidden_layers") + if len(self.mlp_layer_types) != self.num_hidden_layers: + raise ValueError("mlp_layer_types length must equal num_hidden_layers") + if set(self.layer_types) - set(BEAM_LAYER_TYPES): + raise ValueError(f"unsupported layer_types: {sorted(set(self.layer_types))}") + if set(self.mlp_layer_types) - set(BEAM_MLP_LAYER_TYPES): + raise ValueError(f"unsupported mlp_layer_types: {sorted(set(self.mlp_layer_types))}") + if self.hidden_act != "silu": + raise ValueError(f"hidden_act must be 'silu': {self.hidden_act!r}") + if self.num_attention_heads % self.num_key_value_heads: + raise ValueError("num_attention_heads must be divisible by num_key_value_heads") + if "sliding_attention" in self.layer_types and not self.sliding_window: + raise ValueError("sliding_attention layers need a positive sliding_window") + if "sparse" in self.mlp_layer_types: + if self.num_experts % self.n_group: + raise ValueError("num_experts must be divisible by n_group") + if not 1 <= self.topk_group <= self.n_group: + raise ValueError("topk_group must be in [1, n_group]") + if self.num_experts // self.n_group < 2: + raise ValueError("Beam group scores sum the top-2 experts of each group") + if self.num_experts_per_tok > self.topk_group * (self.num_experts // self.n_group): + raise ValueError("num_experts_per_tok exceeds the experts of the selected groups") + + +# --------------------------------------------------------------------------- primitives + + +def beam_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """RMSNorm over the last dim: fp32 statistics and affine, one cast to ``x.dtype``.""" + in_dtype = x.dtype + x32 = x.float() + rms = x32.pow(2).mean(dim=-1, keepdim=True).add(eps).rsqrt() + normed = x32 * rms + return (normed * weight).to(in_dtype) + + +def beam_rope_cos_sin( + positions: torch.Tensor, head_dim: int, rope_theta: float, factor: float +) -> Tuple[torch.Tensor, torch.Tensor]: + """fp32 ``(cos, sin)`` of shape ``positions.shape + (head_dim,)`` (NeoX halves). + + Linear position interpolation divides the inverse frequencies by ``factor``. + """ + exponent = torch.arange(0, head_dim, 2, dtype=torch.float32, device=positions.device) + inv_freq = 1.0 / (rope_theta ** (exponent / head_dim)) + inv_freq = inv_freq / factor + freqs = positions.float().unsqueeze(-1) * inv_freq + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos(), emb.sin() + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def beam_apply_rotary( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """Rotate ``q``/``k`` of shape ``(..., heads, head_dim)`` in fp32, then cast back. + + ``cos``/``sin`` have shape ``(..., head_dim)`` and broadcast over the head axis. + """ + q_dtype, k_dtype = q.dtype, k.dtype + cos = cos.unsqueeze(-2).float() + sin = sin.unsqueeze(-2).float() + q = q.float() + k = k.float() + q_out = q * cos + _rotate_half(q) * sin + k_out = k * cos + _rotate_half(k) * sin + return q_out.to(q_dtype), k_out.to(k_dtype) + + +def beam_ssmax_scale(positions: torch.Tensor) -> torch.Tensor: + """Scalable-Softmax query scale ``exp(0.1 * log1p((p + 1) / 2048))`` in fp32.""" + return torch.exp(0.1 * torch.log1p((positions.float() + 1) / 2048)) + + +def beam_apply_ssmax(q: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Scale ``q`` of shape ``positions.shape + (heads, head_dim)`` in fp32, then cast.""" + scale = beam_ssmax_scale(positions).unsqueeze(-1).unsqueeze(-1) + return (q * scale).to(q.dtype) + + +def beam_output_gate( + attention_output: torch.Tensor, gate: torch.Tensor, gate_bias: torch.Tensor +) -> torch.Tensor: + """Sigmoid output gate on ``(..., heads, head_dim)`` tensors, in the activation dtype.""" + output_dtype = attention_output.dtype + gate = gate + gate_bias.unsqueeze(-1) + return (attention_output * torch.sigmoid(gate)).to(output_dtype) + + +def beam_depth_coefficients(layer_idx: int, is_attention: bool) -> Tuple[float, float]: + """``(x_scale, residual_scale)`` for the attention or FFN add of a zero-based layer.""" + one_based = layer_idx + 1 + num_adds = 2 * one_based if is_attention else 2 * one_based + 1 + denom = math.sqrt(num_adds) + return 1.0 / denom, math.sqrt(num_adds - 1) / denom + + +def beam_depth_residual( + x: torch.Tensor, residual: torch.Tensor, scales: Tuple[float, float] +) -> torch.Tensor: + """``x * x_scale + residual * residual_scale`` in the activation dtype.""" + x_scale, residual_scale = scales + return x * x_scale + residual * residual_scale + + +def beam_softcap(logits: torch.Tensor, cap: float) -> torch.Tensor: + """``tanh(logits / cap) * cap`` computed in fp32.""" + return (logits.float() / cap).tanh() * cap + + +class BeamRMSNorm(nn.Module): + """RMSNorm whose affine weight is stored in the activation dtype.""" + + def __init__(self, dim: int, eps: float, dtype: torch.dtype): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=dtype), requires_grad=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return beam_rms_norm(x, self.weight, self.eps) + + +class BeamSwiGLU(nn.Module): + """Bias-free SwiGLU with the activation evaluated in fp32.""" + + def __init__(self, model_config: ModelConfig, hidden_size: int, intermediate_size: int): + super().__init__() + dtype = model_config.pretrained_config.torch_dtype + quant_config = model_config.get_quant_config() + skip = model_config.skip_create_weights_in_init + self.gate_proj = Linear( + hidden_size, + intermediate_size, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip, + ) + self.up_proj = Linear( + hidden_size, + intermediate_size, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip, + ) + self.down_proj = Linear( + intermediate_size, + hidden_size, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate = self.gate_proj(x) + up = self.up_proj(x) + activated = (F.silu(gate.float()) * up.float()).to(gate.dtype) + return self.down_proj(activated) + + +class BeamRouting(NamedTuple): + """Router outputs and the intermediates that define them (all fp32 except indices).""" + + logits: torch.Tensor + scores: torch.Tensor + group_scores: torch.Tensor + top_groups: torch.Tensor + excluded: torch.Tensor + indices: torch.Tensor + weights: torch.Tensor + + +class BeamRouter(nn.Module): + """Sigmoid grouped top-k router; the correction bias only affects selection. + + Selection: sigmoid scores plus the correction bias, group score = sum of the top-2 + corrected scores in each group, sorted top-``topk_group`` groups, then sorted + top-``top_k`` experts among them (other groups masked to ``-inf``). Weights are + the unbiased scores of the selected experts in selection order, optionally + normalized by their sum plus 1e-20; the routed scaling factor is not applied here. + """ + + def __init__(self, config: BeamConfig): + super().__init__() + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_tok + self.n_group = config.n_group + self.topk_group = config.topk_group + self.norm_topk_prob = config.norm_topk_prob + self.weight = nn.Parameter( + torch.empty(config.num_experts, config.hidden_size, dtype=torch.float32), + requires_grad=False, + ) + self.register_buffer( + "e_score_correction_bias", torch.zeros(config.num_experts, dtype=torch.float32) + ) + + def route(self, hidden_states: torch.Tensor) -> BeamRouting: + """Route ``(N, hidden)`` tokens; graph-capturable (no host synchronization).""" + logits = F.linear(hidden_states.float(), self.weight.float()) + scores = logits.sigmoid() + scores_for_choice = scores + self.e_score_correction_bias + num_tokens, num_experts = scores.shape + experts_per_group = num_experts // self.n_group + grouped = scores_for_choice.view(num_tokens, self.n_group, experts_per_group) + group_scores = grouped.topk(2, dim=-1).values.sum(dim=-1) + top_groups = group_scores.topk(self.topk_group, dim=-1, sorted=True).indices + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, top_groups, 1) + expert_mask = ( + group_mask.unsqueeze(-1) + .expand(num_tokens, self.n_group, experts_per_group) + .reshape(num_tokens, num_experts) + ) + excluded = ~expert_mask.bool() + scores_for_choice = scores_for_choice.masked_fill(excluded, float("-inf")) + indices = scores_for_choice.topk(self.top_k, dim=-1, sorted=True).indices + weights = scores.gather(-1, indices) + if self.norm_topk_prob: + weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20) + return BeamRouting(logits, scores, group_scores, top_groups, excluded, indices, weights) + + def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + routing = self.route(hidden_states) + return routing.weights, routing.indices + + +class BeamMoE(nn.Module): + """Routed experts with the reference combine, executed with static shapes. + + Every expert evaluates every token row in ascending expert order; a route + contributes only where the router selected that expert. Each selected route is + ``expert(x) * weight`` rounded to the activation dtype and accumulated into an + fp32 sum, which is scaled by ``routed_scaling_factor`` and cast once. The loop has + no data-dependent branch or host synchronization, so it runs inside CUDA graphs. + """ + + def __init__(self, model_config: ModelConfig): + super().__init__() + config = model_config.pretrained_config + self.routed_scaling_factor = config.routed_scaling_factor + self.gate = BeamRouter(config) + self.experts = nn.ModuleList( + [ + BeamSwiGLU(model_config, config.hidden_size, config.moe_intermediate_size) + for _ in range(config.num_experts) + ] + ) + + def combine( + self, hidden_states: torch.Tensor, weights: torch.Tensor, indices: torch.Tensor + ) -> torch.Tensor: + """Expert outputs for given routing: ``hidden_states`` is ``(N, hidden)``.""" + out = torch.zeros(hidden_states.shape, dtype=torch.float32, device=hidden_states.device) + for expert_id, expert in enumerate(self.experts): + selected = indices == expert_id + route_weight = (weights * selected).sum(dim=-1, keepdim=True) + route = (expert(hidden_states) * route_weight).to(hidden_states.dtype) + hit = selected.any(dim=-1, keepdim=True) + out = torch.where(hit, out + route.float(), out) + out = out * self.routed_scaling_factor + return out.to(hidden_states.dtype) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + weights, indices = self.gate(hidden_states) + return self.combine(hidden_states, weights, indices) + + +# ---------------------------------------------------------------------------- attention + + +class BeamAttention(Attention): + """Beam GQA attention on the standard module/backend/KV-cache path. + + Q/K RMSNorm, RoPE (sliding layers) or SSMax (full layers) and the gate bias run in + the module before the backend appends K/V, so the cache holds post-norm, + post-RoPE keys. Backend RoPE fusion is disabled. + """ + + def __init__(self, model_config: ModelConfig[BeamConfig], layer_idx: int): + config = model_config.pretrained_config + self.is_sliding = config.layer_types[layer_idx] == "sliding_attention" + self.attention_window_size = config.sliding_window if self.is_sliding else None + rope = config.rope_parameters["sliding_attention"] + self.rope_theta = float(rope["rope_theta"]) + self.rope_factor = float(rope.get("factor", 1.0)) + super().__init__( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + max_position_embeddings=config.max_position_embeddings, + bias=config.attention_bias, + pos_embd_params=None, + rope_fusion=False, + layer_idx=layer_idx, + dtype=config.torch_dtype, + # The output-projection bias is added after the bias-free GEMM. + dense_bias=False, + config=model_config, + attn_output_gate=True, + head_dim=config.head_dim, + ) + dtype = config.torch_dtype + self.gate_bias = nn.Parameter(torch.zeros(self.num_heads, dtype=dtype), requires_grad=False) + self.o_proj_bias = ( + nn.Parameter(torch.zeros(config.hidden_size, dtype=dtype), requires_grad=False) + if config.attention_bias + else None + ) + self.q_norm = BeamRMSNorm(self.head_dim, config.rms_norm_eps, dtype) + self.k_norm = BeamRMSNorm(self.head_dim, config.rms_norm_eps, dtype) + + def apply_rope( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + position_ids: torch.Tensor, + ): + """Q/K norm, then RoPE (sliding) or SSMax (full) at the tokens' true positions.""" + q, k, v = self.split_qkv(q, k, v) + if position_ids is None: + raise ValueError("Beam attention needs per-token position ids") + num_tokens = q.shape[0] + positions = position_ids.reshape(-1)[:num_tokens] + q = self.q_norm(q.reshape(num_tokens, self.num_heads, self.head_dim)) + k = self.k_norm(k.reshape(num_tokens, self.num_key_value_heads, self.head_dim)) + if self.is_sliding: + cos, sin = beam_rope_cos_sin( + positions, self.head_dim, self.rope_theta, self.rope_factor + ) + q, k = beam_apply_rotary(q, k, cos, sin) + else: + q = beam_apply_ssmax(q, positions) + return (q.reshape(num_tokens, self.q_size), k.reshape(num_tokens, self.kv_size), v) + + def apply_output_gate(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + num_tokens = attention_output.shape[0] + out = attention_output.reshape(num_tokens, self.num_heads, self.head_dim) + gate = gate[:num_tokens].reshape(num_tokens, self.num_heads, self.head_dim) + gated = beam_output_gate(out, gate, self.gate_bias) + return gated.reshape(num_tokens, self.num_heads * self.head_dim) + + def forward( + self, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + kwargs.setdefault("attention_window_size", self.attention_window_size) + output = super().forward( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + **kwargs, + ) + if self.o_proj_bias is not None: + output = output + self.o_proj_bias + return output + + +# ------------------------------------------------------------------------------- blocks + + +class BeamDecoderLayer(DecoderLayer): + """Sandwich-norm decoder block with depth-scaled residual adds.""" + + def __init__(self, model_config: ModelConfig[BeamConfig], layer_idx: int): + super().__init__() + config = model_config.pretrained_config + dtype = config.torch_dtype + eps = config.rms_norm_eps + self.layer_idx = layer_idx + self.input_layernorm = BeamRMSNorm(config.hidden_size, eps, dtype) + self.self_attn = BeamAttention(model_config, layer_idx) + self.post_attention_layernorm = BeamRMSNorm(config.hidden_size, eps, dtype) + self.pre_feedforward_layernorm = BeamRMSNorm(config.hidden_size, eps, dtype) + self.is_sparse = config.mlp_layer_types[layer_idx] == "sparse" + if self.is_sparse: + self.mlp = BeamMoE(model_config) + else: + self.mlp = BeamSwiGLU(model_config, config.hidden_size, config.intermediate_size) + self.post_feedforward_layernorm = BeamRMSNorm(config.hidden_size, eps, dtype) + self.attn_scales = beam_depth_coefficients(layer_idx, is_attention=True) + self.mlp_scales = beam_depth_coefficients(layer_idx, is_attention=False) + + def forward( + self, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + residual: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + residual = hidden_states + x = self.input_layernorm(hidden_states) + x = self.self_attn( + position_ids=position_ids, hidden_states=x, attn_metadata=attn_metadata, **kwargs + ) + x = self.post_attention_layernorm(x) + hidden_states = beam_depth_residual(x, residual, self.attn_scales) + residual = hidden_states + x = self.pre_feedforward_layernorm(hidden_states) + x = self.mlp(x) + x = self.post_feedforward_layernorm(x) + return beam_depth_residual(x, residual, self.mlp_scales) + + +class BeamModel(DecoderModel): + def __init__(self, model_config: ModelConfig[BeamConfig]): + super().__init__(model_config) + config = model_config.pretrained_config + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + mapping=model_config.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + ) + self.layers = nn.ModuleList( + [ + BeamDecoderLayer(model_config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = BeamRMSNorm(config.hidden_size, config.rms_norm_eps, config.torch_dtype) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("Specify exactly one of input_ids and inputs_embeds") + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + for layer in self.layers: + hidden_states = layer( + position_ids=position_ids, hidden_states=hidden_states, attn_metadata=attn_metadata + ) + return self.norm(hidden_states) + + +class BeamForCausalLM(DecoderModelForCausalLM[BeamModel, BeamConfig]): + """Beam causal LM: last-token (or context) logits in fp32, softcapped when enabled.""" + + def __init__(self, model_config: ModelConfig[BeamConfig]): + super().__init__( + BeamModel(model_config), + config=model_config, + hidden_size=model_config.pretrained_config.hidden_size, + vocab_size=model_config.pretrained_config.vocab_size, + ) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + **kwargs, + ) -> torch.Tensor: + logits = super().forward( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + return_context_logits=return_context_logits, + **kwargs, + ) + cap = self.config.final_logit_softcapping + if cap is not None: + logits = beam_softcap(logits, cap) + return logits + + def load_weights( + self, + weights: Dict[str, torch.Tensor], + weight_mapper=None, + skip_modules: Optional[List[str]] = None, + params_map: Optional[Dict[str, str]] = None, + allow_partial_loading: bool = False, + ) -> None: + """Load a reference-layout state dict exactly (values and dtypes, no casts).""" + if weight_mapper is not None or params_map or skip_modules or allow_partial_loading: + raise ValueError("BeamForCausalLM loads complete reference-layout state dicts only") + if self.model_config.mapping.tp_size > 1: + raise ValueError( + "BeamForCausalLM loads unsharded state dicts; use tp_size=1 or the " + "ModelingV2 TP4 target" + ) + load_beam_state_dict(self, weights) + + +# ------------------------------------------------------------------------ weight mapping + + +def beam_parameter_dtype(name: str, activation_dtype: torch.dtype) -> torch.dtype: + """Source dtype policy: router weight and correction bias fp32, all else activation.""" + return torch.float32 if name.endswith(BEAM_FP32_PARAMETER_SUFFIXES) else activation_dtype + + +def pack_beam_qkv_gate( + q: torch.Tensor, + gate: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + num_heads: int, + head_dim: int, +) -> torch.Tensor: + """Fused projection rows: per head ``[q_h, gate_h]``, then all K rows, then all V rows.""" + trailing = q.shape[1:] + q_gate = torch.stack( + (q.reshape(num_heads, head_dim, *trailing), gate.reshape(num_heads, head_dim, *trailing)), + dim=1, + ) + return torch.cat((q_gate.reshape(2 * num_heads * head_dim, *trailing), k, v), dim=0) + + +def unpack_beam_qkv_gate( + fused: torch.Tensor, num_heads: int, num_kv_heads: int, head_dim: int +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Inverse of :func:`pack_beam_qkv_gate`: returns ``(q, gate, k, v)``.""" + q_width = num_heads * head_dim + kv_width = num_kv_heads * head_dim + q_gate, k, v = fused.split([2 * q_width, kv_width, kv_width], dim=0) + q_gate = q_gate.reshape(num_heads, 2, head_dim, *fused.shape[1:]) + q = q_gate[:, 0].reshape(q_width, *fused.shape[1:]) + gate = q_gate[:, 1].reshape(q_width, *fused.shape[1:]) + return q, gate, k, v + + +def map_beam_attention_state( + state: Dict[str, torch.Tensor], config: BeamConfig, prefix: str = "" +) -> Dict[str, torch.Tensor]: + """Map one attention module's reference entries (under ``prefix``) to candidate names. + + The separate Q, gate, K and V projections become the fused ``qkv_proj``. With + ``attention_bias`` the gate rows of the fused bias are zero (the gate projection has + no bias) and the output-projection bias becomes ``o_proj_bias``. Consumed entries are + removed from ``state``. + """ + heads, head_dim = config.num_attention_heads, config.head_dim + mapped: Dict[str, torch.Tensor] = {} + q = state.pop(prefix + "q_proj.weight") + gate = state.pop(prefix + "gate_proj.weight") + k = state.pop(prefix + "k_proj.weight") + v = state.pop(prefix + "v_proj.weight") + mapped[prefix + "qkv_proj.weight"] = pack_beam_qkv_gate(q, gate, k, v, heads, head_dim) + if config.attention_bias: + qb = state.pop(prefix + "q_proj.bias") + kb = state.pop(prefix + "k_proj.bias") + vb = state.pop(prefix + "v_proj.bias") + mapped[prefix + "qkv_proj.bias"] = pack_beam_qkv_gate( + qb, torch.zeros_like(qb), kb, vb, heads, head_dim + ) + mapped[prefix + "o_proj_bias"] = state.pop(prefix + "o_proj.bias") + mapped[prefix + "o_proj.weight"] = state.pop(prefix + "o_proj.weight") + for name in ("gate_bias", "q_norm.weight", "k_norm.weight"): + mapped[prefix + name] = state.pop(prefix + name) + return mapped + + +def map_beam_state_dict( + state: Dict[str, torch.Tensor], + config: BeamConfig, + attention_prefixes: Optional[List[str]] = None, + tied_lm_head: bool = False, +) -> Dict[str, torch.Tensor]: + """Map reflection-core state-dict names to candidate parameter/buffer names. + + ``attention_prefixes`` lists the attention modules contained in ``state`` (default: + every ``model.layers.{i}.self_attn.`` of the full model). Other entries keep their + names. With ``tied_lm_head`` a supplied ``lm_head.weight`` must equal the embedding. + """ + remaining = dict(state) + if attention_prefixes is None: + attention_prefixes = [ + f"model.layers.{i}.self_attn." for i in range(config.num_hidden_layers) + ] + mapped: Dict[str, torch.Tensor] = {} + for prefix in attention_prefixes: + mapped.update(map_beam_attention_state(remaining, config, prefix)) + if tied_lm_head: + tied = remaining.pop("lm_head.weight", None) + if tied is not None and not torch.equal(tied, state["model.embed_tokens.weight"]): + raise ValueError("tied lm_head.weight differs from model.embed_tokens.weight") + mapped.update(remaining) + return mapped + + +def beam_state_targets(module: nn.Module, tied_lm_head: bool = False) -> Dict[str, torch.Tensor]: + """Every loadable tensor of ``module``: all parameters plus router correction buffers. + + With ``tied_lm_head`` the LM-head weight is the embedding parameter and appears once. + """ + targets: Dict[str, torch.Tensor] = dict(module.named_parameters()) + for name, buffer in module.named_buffers(): + if name.endswith("e_score_correction_bias"): + targets[name] = buffer + if tied_lm_head: + targets.pop("lm_head.weight", None) + return targets + + +def load_beam_state_dict( + module: nn.Module, + state: Dict[str, torch.Tensor], + config: Optional[BeamConfig] = None, + attention_prefixes: Optional[List[str]] = None, +) -> Dict[str, str]: + """Copy a reference-layout state dict into ``module``; return each tensor's dtype. + + ``module`` is a ``BeamForCausalLM`` (default prefixes) or a standalone Beam module + with explicit ``config`` and ``attention_prefixes`` (``["self_attn."]`` for a + decoder layer, ``[""]`` for an attention module, ``[]`` for none). Requires an exact + name match after mapping, identical shapes and the source dtype policy on both + sides (no implicit casts). Returns the loaded dtypes for accounting. + """ + is_model = config is None + config = config if config is not None else module.config + tied = is_model and bool(config.tie_word_embeddings) + mapped = map_beam_state_dict(state, config, attention_prefixes, tied_lm_head=tied) + targets = beam_state_targets(module, tied_lm_head=tied) + missing = sorted(set(targets) - set(mapped)) + unexpected = sorted(set(mapped) - set(targets)) + if missing or unexpected: + raise ValueError(f"state dict mismatch: missing={missing} unexpected={unexpected}") + accounting: Dict[str, str] = {} + activation_dtype = config.torch_dtype + with torch.no_grad(): + for name, target in targets.items(): + source = mapped[name] + expected = beam_parameter_dtype(name, activation_dtype) + if source.dtype != expected or target.dtype != expected: + raise ValueError( + f"{name}: source {source.dtype} / target {target.dtype}, policy {expected}" + ) + if source.shape != target.shape: + raise ValueError(f"{name}: shape {tuple(source.shape)} != {tuple(target.shape)}") + target.copy_(source) + accounting[name] = str(target.dtype) + return accounting diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 14af68286a0c..1758cae0452b 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # from dataclasses import replace -from typing import Dict, Optional +from typing import Dict, NamedTuple, Optional import torch import torch.nn.functional as F @@ -14,6 +14,7 @@ from ...functional import RotaryScalingType from ..attention.rotary_embedding import RotaryEmbedding +from ..modules.logits_processor import LogitsProcessor try: from ..custom_ops import flashinfer_apply_rope_with_cos_sin_cache_inplace as _flashinfer_rope @@ -169,6 +170,127 @@ def dflash2_walk_candidate_paths( return realized +def dflash_final_logit_softcap(pretrained_config: PretrainedConfig) -> Optional[float]: + """The drafter checkpoint's ``dflash_config.final_logit_softcapping``, or None if absent.""" + dflash_config = getattr(pretrained_config, "dflash_config", None) or {} + cap = dflash_config.get("final_logit_softcapping") + if cap is None: + return None + if isinstance(cap, bool) or not isinstance(cap, (int, float)) or not cap > 0: + raise ValueError( + f"dflash_config.final_logit_softcapping must be a positive number, got {cap!r}" + ) + return float(cap) + + +class DFlashSoftcapLogitsProcessor(LogitsProcessor): + """Draft logits of a drafter whose checkpoint declares ``final_logit_softcapping``. + + The native processor's projection through the draft head and its FP32 widening, then + ``cap * tanh(logits / cap)`` exactly once. The DFlash worker calls the drafter's + ``logits_processor`` before its global candidate merge, so with a vocabulary-sharded head + each TP rank caps its own shard; the candidate selector's unary scores are those capped + values, and nothing caps again after selection. + """ + + def __init__(self, cap: float): + super().__init__() + if isinstance(cap, bool) or not isinstance(cap, (int, float)) or not cap > 0: + raise ValueError(f"the draft logit softcap must be a positive number, got {cap!r}") + self.cap = float(cap) + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: nn.Module, + attn_metadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + logits = super().forward(hidden_states, lm_head, attn_metadata, return_context_logits) + return self.cap * torch.tanh(logits / self.cap) + + +class DFlashWindowGather(NamedTuple): + """Where one draft block's non-causal windowed attention reads its keys, per batch row. + + ``positions`` ``[B, S]`` are the slot-local key positions the block's queries can reach, + ``S = window + block_size - 1`` (see :func:`dflash_window_gather`); ``valid`` marks those + that hold a key of the request (its committed context and the current block, nothing at + or past ``num_ctx + block_size``); ``mask`` ``[B, 1, block_size, S]`` is query ``j``'s + visibility: a valid key at most ``window - 1`` positions from ``num_ctx + j``. ``pages`` + and ``offsets`` address every position in the paged pool; an invalid position is + redirected to the block's own first slot, a page the request owns and the append wrote. + """ + + positions: torch.Tensor + valid: torch.Tensor + mask: torch.Tensor + pages: torch.Tensor + offsets: torch.Tensor + + +def dflash_window_gather( + block_tables: torch.Tensor, + num_ctx: torch.Tensor, + block_size: int, + window: int, + page_size: int, +) -> DFlashWindowGather: + """Key addressing of a non-causal ``window``-token block attention over a paged pool. + + Query ``j`` of a request's block sits at slot-local position ``num_ctx + j`` (the block's + K/V were appended at ``num_ctx .. num_ctx + block_size - 1``) and sees every key of the + request within ``window - 1`` positions on either side, so the block together spans + ``[num_ctx - (window - 1), num_ctx + block_size)``. A single most-recent-``window`` gather + would be wrong for every query but the last. + + Args: + block_tables: ``[B, pages_per_row]`` page indices of each batch row's slot-local pages. + num_ctx: ``[B]`` context length per row before the block. + Fixed shapes and device arithmetic only: no host synchronization. + """ + device = num_ctx.device + ctx = num_ctx.long() + span = window + block_size - 1 + positions = (ctx - (window - 1)).unsqueeze(1) + torch.arange(span, device=device) + valid = (positions >= 0) & (positions < (ctx + block_size).unsqueeze(1)) + queries = ctx.unsqueeze(1) + torch.arange(block_size, device=device) + distance = (queries.unsqueeze(2) - positions.unsqueeze(1)).abs() + mask = (valid.unsqueeze(1) & (distance <= window - 1)).unsqueeze(1) + safe = torch.where(valid, positions, ctx.unsqueeze(1)) + pages = block_tables.long().gather(1, torch.div(safe, page_size, rounding_mode="floor")) + return DFlashWindowGather(positions, valid, mask, pages, safe % page_size) + + +def dflash_windowed_block_attention( + q: torch.Tensor, + layer_cache: torch.Tensor, + gather: DFlashWindowGather, + scale: float, +) -> torch.Tensor: + """Non-causal windowed attention of draft blocks over the HND paged pool, in torch SDPA. + + ``q`` ``[B, block_size, Hq, D]`` (post-norm/RoPE), ``layer_cache`` one layer's pool + ``[pages, 2, Hkv, page_size, D]`` holding the context and the block's own K/V. Gathers the + ``S`` positions ``gather`` names, zeroes the invalid ones and lets ``gather.mask`` select + each query's keys; Hq / Hkv query heads share a KV head. Returns ``[B, block_size, Hq, D]``. + """ + keep = gather.valid.unsqueeze(-1).unsqueeze(-1) + k = layer_cache[gather.pages, 0, :, gather.offsets].masked_fill(~keep, 0) + v = layer_cache[gather.pages, 1, :, gather.offsets].masked_fill(~keep, 0) + out = F.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + attn_mask=gather.mask, + dropout_p=0.0, + is_causal=False, + scale=scale, + enable_gqa=True, + ) + return out.transpose(1, 2) + + class DFlash2BlockConv(nn.Module): """Two-tap dynamic depthwise convolutions wrapped around one sublayer. @@ -1306,12 +1428,13 @@ def _resolve_block_attention(self, layer_idx: int) -> tuple[bool, tuple[int, int return causal, window_size def validate_block_attention_windows(self) -> None: - """Refuse windows the TRTLLM block-decode kernel cannot apply. + """Refuse windows the TRTLLM block decode cannot apply. TRTLLM-Gen only implements sliding windows for causal masks: flashinfer - rejects ``causal=False`` with any finite ``window_left`` outright. A - non-causal windowed layer must therefore run on VANILLA or FA4, which - take a two-sided window. + rejects ``causal=False`` with any finite ``window_left`` outright. On the + TRTLLM backend a non-causal layer with a two-sided window of equal sides + runs :func:`dflash_windowed_block_attention` over the same paged pool + instead; any other non-causal windowed layer must run on VANILLA or FA4. """ if self.dflash_attention_backend != "TRTLLM": return @@ -1319,14 +1442,15 @@ def validate_block_attention_windows(self) -> None: if num_layers is None: num_layers = len(self.model.layers) for layer_idx in range(num_layers): - causal, (window_left, _) = self._resolve_block_attention(layer_idx) - if causal or window_left < 0: + causal, (window_left, window_right) = self._resolve_block_attention(layer_idx) + if causal or window_left < 0 or window_left == window_right: continue raise ValueError( f"DFlash draft layer {layer_idx} attends non-causally within a " - f"{window_left + 1}-token sliding window. The TRTLLM DFlash " - "attention backend does not support non-causal sliding-window " - "attention. Use attention_backend=VANILLA or FA4." + f"window of {window_left} keys before and {window_right} after " + "each query. The TRTLLM DFlash attention backend supports a " + "non-causal window only with equal sides. Use " + "attention_backend=VANILLA or FA4." ) def _prepare_dflash_trtllm_gen_buffers( @@ -1580,6 +1704,8 @@ def dflash_forward( residual = None has_block_conv = self.has_block_conv + # Non-causal windowed layers share one key addressing per window size. + window_gathers: dict[int, DFlashWindowGather] = {} for layer_idx, layer in enumerate(self.model.layers): attn_mod = layer.self_attn @@ -1721,6 +1847,19 @@ def dflash_forward( bmm1_scale_log2=None, multi_ctas_kv_counter_buffer=self._dflash_trtllm_gen_counters, ) + elif window_left >= 0: + # TRTLLM-Gen has no non-causal sliding window (see + # validate_block_attention_windows): torch SDPA over the + # same pages the append just wrote. + gather = window_gathers.get(window_left) + if gather is None: + gather = dflash_window_gather( + block_tables, cache_seqlens_i32, block_size, window_left + 1, page_size + ) + window_gathers[window_left] = gather + out = dflash_windowed_block_attention( + Q_bshd, layer_cache, gather, head_dim**-0.5 + ) else: cum_seq_lens_q = torch.arange( 0, diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 24b094deb4d4..b3a698804e08 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -1117,6 +1117,13 @@ def load_pretrained_config(model_name_or_path: str, config_class = _CONFIG_REGISTRY[model_type] model_config = config_class.from_pretrained(model_name_or_path, **kwargs) + elif model_type == "beam": + # Beam is not a transformers model type; its config class lives with + # the native model. The checkpoint directory is kept as name_or_path, + # as from_pretrained does, for targets that read the stored config. + from tensorrt_llm._torch.models.modeling_beam import BeamConfig + model_config = BeamConfig.from_dict(config_dict, **kwargs) + model_config.name_or_path = model_name_or_path elif model_type == "minicpmv4_6" or ( architectures and architectures[0] == "MiniCPMV4_6ForConditionalGeneration"): diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 3c10a654ba35..c97cdc6cbbf0 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -137,6 +137,10 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] # TIMEOUT measured: 3122 cases in 13m54s on one GB300 (the cpu_only files run on the CPU # stage instead), so 30 is 2.2x headroom. + # Beam's generated-input operator, block, six-layer and metadata regressions are + # single-GPU L0 coverage, including the high_cuda_memory cases. Real-checkpoint + # TP4 smoke coverage belongs in a separate nightly/LR job. Full Beam GSM8K + # evaluation is manual-only, excluded from L0, nightly and LR CI. - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30) - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] # qwen4_exp target, pre-merge: gsm8k; one MTP engine for gsm8k with the acceptance-length diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index dce2d3668c00..3b246cc4a0ff 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -66,6 +66,8 @@ l0_cpu: # The rest of unittest/_torch/modeling_v2 needs a GPU. - unittest/_torch/modeling_v2/test_modeling_v2_claims.py - unittest/_torch/modeling_v2/test_modeling_v2_routing.py + - unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_serving.py + - unittest/_torch/modeling_v2/models/beam/test_loading_contracts.py - unittest/_torch/modeling_v2/test_modeling_v2_deepseek_default_shards.py - unittest/_torch/modeling_v2/test_modeling_v2_gpt_oss_default_shards.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_dense_backend_library_order diff --git a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py index fac6b0eb5359..208e937d4d40 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py +++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py @@ -134,6 +134,54 @@ def test_small_context_fallback_preserves_mixed_batch_generation( assert fallback.is_supported(q, None, None, metadata, forward_args, phase=phase) +@pytest.mark.cpu_only +@pytest.mark.parametrize("fallback_has_kernel", [True, False]) +def test_small_paged_context_uses_fallback_only_with_a_fused_kernel( + monkeypatch: pytest.MonkeyPatch, fallback_has_kernel: bool +) -> None: + """A paged (cached-KV) small BF16 context leaves FlashInfer only when the fallback + build has a fused paged-context kernel for its head size and page size.""" + from tensorrt_llm._torch.attention.backends.fmha import flashinfer_trtllm_gen + + lookups = [] + + def kernel_exists(head_size, kv_cache_dtype, tokens_per_block): + lookups.append((head_size, kv_cache_dtype, tokens_per_block)) + return fallback_has_kernel + + monkeypatch.setattr(flashinfer_trtllm_gen, "_fallback_has_paged_context_kernel", kernel_exists) + monkeypatch.setattr(flashinfer_trtllm_gen, "get_sm_version", lambda: 100) + attn = FakeAttention() + attn.head_dim = 128 + attn.sparse_params = None + attn.position_embedding_type = 0 + fmha = FlashInferTrtllmGenFmha(attn) + metadata = SimpleNamespace( + num_contexts=1, + num_generations=0, + helix_position_offsets=None, + num_sparse_topk=0, + use_spec_decoding=False, + kv_cache_block_offsets=object(), + kv_cache_manager=None, + tokens_per_block=64, + is_cross=False, + beam_width=1, + use_paged_context_fmha=True, + ) + q = torch.empty((256, 3 * attn.head_dim), dtype=torch.bfloat16) + forward_args = AttentionForwardArgs( + output=torch.empty((256, attn.head_dim), dtype=torch.bfloat16), + attention_input_type=AttentionInputType.context_only, + is_fused_qkv=True, + ) + supported, reason = fmha._is_supported_with_reason( + q, None, None, attn, metadata, forward_args, phase=FmhaPhase.CONTEXT + ) + assert supported is not fallback_has_kernel, reason + assert lookups == [(128, DataType.BF16, 64)] + + @pytest.mark.parametrize("kv_cache_dtype", [DataType.FP8, DataType.NVFP4]) @pytest.mark.parametrize( "attention_mask", [PredefinedAttentionMask.CAUSAL, CustomAttentionMask.CUSTOM] diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_neox_rope.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_neox_rope.py new file mode 100644 index 000000000000..3eef5075bb85 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_neox_rope.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_neox_rope entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_neox_rope import ( + beam_neox_rope as op, +) + +__extra_import_path__ = [".."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_neox_rope requires a CUDA device" + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + t, d, dtype = spec["tokens"], spec["head_dim"], spec["dtype"] + q = torch.randn(t, spec["num_heads"], d, generator=g, device="cuda").to(dtype) + k = torch.randn(t, spec["num_kv_heads"], d, generator=g, device="cuda").to(dtype) + return q, k, torch.arange(t, device="cuda") + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bitwise: the op and its reference make the same fp32 operations in the same order + and round once, so any difference is a changed formula, not accumulation noise.""" + certifying(op) + spec = cell.spec + q, k, positions = _build(spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + q_out, k_out = op(q, k, positions, spec["theta"], spec["factor"]) + q_ref, k_ref = op.reference(q, k, positions, spec["theta"], spec["factor"]) + for out, ref, src in ((q_out, q_ref, q), (k_out, k_ref, k)): + assert out.shape == src.shape and out.dtype == src.dtype + assert out.data_ptr() != src.data_ptr(), "the op must return a new tensor" + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) + if spec["tokens"] == 1: + assert torch.equal(q_out, q) and torch.equal(k_out, k), "position 0 must rotate by 0" + else: + assert not torch.equal(q_out[1:], q[1:]), "positions past 0 must rotate" + + +def test_halves_rotate_together_not_interleaved_pairs() -> None: + """A unit vector in dimension 0 at position p goes to (cos p, ..., sin p at D/2, ...): + NeoX pairs dimension i with i + D/2, not with i + 1.""" + d = 128 + q = torch.zeros(2, 1, d, dtype=torch.float32, device="cuda") + q[:, 0, 0] = 1.0 + k = q.clone() + positions = torch.arange(2, device="cuda") + q_out, _ = op(q, k, positions, 10000.0, 1.0) + assert q_out[1, 0, 0].item() == pytest.approx(torch.cos(torch.tensor(1.0)).item(), abs=1e-6) + assert q_out[1, 0, d // 2].item() == pytest.approx( + torch.sin(torch.tensor(1.0)).item(), abs=1e-6 + ) + assert q_out[1, 0, 1].item() == 0.0 + + +@pytest.mark.parametrize( + "case,match", + [ + ("positions_broadcast", "positions must be"), + ("float_positions", "integer token positions"), + ("odd_head_dim", "even head_dim"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + q = torch.randn(17, 48, 128, dtype=torch.bfloat16, device="cuda") + k = torch.randn(17, 8, 128, dtype=torch.bfloat16, device="cuda") + positions = torch.arange(17, device="cuda") + if case == "positions_broadcast": + positions = positions[:1] + elif case == "float_positions": + positions = positions.float() + else: + q, k = q[..., :127], k[..., :127] + with pytest.raises(AssertionError, match=match): + with validating(op): + op(q, k, positions, 10000.0, 1.0) diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_nope_query_scale.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_nope_query_scale.py new file mode 100644 index 000000000000..252d5f0a2486 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_nope_query_scale.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_nope_query_scale entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_nope_query_scale import ( + beam_nope_query_scale as op, +) + +__extra_import_path__ = [".."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_nope_query_scale requires a CUDA device" + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + t, h, d, dtype = spec["tokens"], spec["num_heads"], spec["head_dim"], spec["dtype"] + q = torch.randn(t, h, d, generator=g, device="cuda").to(dtype) + positions = torch.arange(t, device="cuda") + return q, positions + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bitwise: the fp32 scale and the single rounding of the product are what the + reference spells out step by step.""" + certifying(op) + q, positions = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(q, positions) + assert out.shape == q.shape and out.dtype == q.dtype + torch.testing.assert_close(out, op.reference(q, positions), rtol=0.0, atol=0.0) + + +def test_scale_against_float64() -> None: + """The fp32 scale is within four fp32 units of the float64 expression at every position, so + the bf16 result differs from a float64-scaled one only where the product sits within that + distance of a bf16 rounding boundary.""" + positions = torch.arange(129, device="cuda") + q = torch.ones(129, 1, 1, dtype=torch.float32, device="cuda") + scale32 = op(q, positions).view(-1) + scale64 = torch.exp(0.1 * torch.log1p((positions.double() + 1) / 2048)) + assert float(((scale32.double() - scale64).abs() / scale64).max()) <= 4 * 2.0**-23 + + +def test_scale_is_visible_at_nonzero_positions() -> None: + """Values one bf16 unit below a power of two round up under the scale once it exceeds half + a unit -- at the later positions of a 129-token context -- while position 0's scale + leaves them unchanged: an omitted scale is detectable only at nonzero positions.""" + positions = torch.arange(129, device="cuda") + q = torch.full((129, 1, 1), 1.9921875, dtype=torch.bfloat16, device="cuda") + out = op(q, positions).view(-1) + assert float(out[0]) == 1.9921875 + assert float(out[-1]) == 2.0 + + +def test_omitting_the_scale_or_rotating_differs() -> None: + """A full-attention q that is left unscaled is a different tensor on ordinary inputs.""" + q, positions = _build(dict(tokens=129, num_heads=48, head_dim=128, dtype=torch.bfloat16), 3) + assert not torch.equal(op(q, positions), q) + + +@pytest.mark.parametrize( + "case,match", + [ + ("positions_length_one", "positions must be"), + ("float_positions", "integer token positions"), + ("q_rank", "q must be"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + q, positions = _build(dict(tokens=17, num_heads=48, head_dim=128, dtype=torch.bfloat16), 11) + if case == "positions_length_one": + positions = positions[:1] + elif case == "float_positions": + positions = positions.float() + else: + q = q.view(17, 48 * 128) + with pytest.raises(AssertionError, match=match): + with validating(op): + op(q, positions) diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_output_gate.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_output_gate.py new file mode 100644 index 000000000000..4a68cc0433bf --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_output_gate.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_output_gate entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_output_gate import ( + beam_output_gate as op, +) + +__extra_import_path__ = [".."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_output_gate requires a CUDA device" + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + t, h, d, dtype = spec["tokens"], spec["num_heads"], spec["head_dim"], spec["dtype"] + attention_output = torch.randn(t, h, d, generator=g, device="cuda").to(dtype) + gate = (2.0 * torch.randn(t, h, d, generator=g, device="cuda")).to(dtype) + gate_bias = (0.75 * torch.randn(h, generator=g, device="cuda")).to(dtype) + return attention_output, gate, gate_bias + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bitwise: each activation-dtype step of the op is an fp32 computation rounded once, + which is what the reference spells out.""" + certifying(op) + attention_output, gate, gate_bias = _build( + cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + ) + out = op(attention_output, gate, gate_bias) + assert out.shape == attention_output.shape and out.dtype == attention_output.dtype + torch.testing.assert_close( + out, op.reference(attention_output, gate, gate_bias), rtol=0.0, atol=0.0 + ) + + +def test_rounding_between_steps_is_observable() -> None: + """One fp32 chain rounded once is a different answer on these inputs, so the cell + gate above distinguishes the source's per-step rounding from a fused gate.""" + attention_output, gate, gate_bias = _build( + dict(tokens=129, num_heads=48, head_dim=128, dtype=torch.bfloat16), seed=7 + ) + fused = ( + attention_output.float() * torch.sigmoid(gate.float() + gate_bias.float().unsqueeze(-1)) + ).to(torch.bfloat16) + assert not torch.equal(op(attention_output, gate, gate_bias), fused) + + +@pytest.mark.parametrize( + "case,match", + [ + ("bias_length_one", "gate_bias must be"), + ("fp32_bias", "share one dtype"), + ("gate_shape", "must both be"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + attention_output, gate, gate_bias = _build( + dict(tokens=17, num_heads=48, head_dim=128, dtype=torch.bfloat16), seed=11 + ) + if case == "bias_length_one": + gate_bias = gate_bias[:1] + elif case == "fp32_bias": + gate_bias = gate_bias.float() + else: + gate = gate[:, :1] + with pytest.raises(AssertionError, match=match): + with validating(op): + op(attention_output, gate, gate_bias) diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_dflash_attention.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_dflash_attention.py new file mode 100644 index 000000000000..add0aa60a906 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_dflash_attention.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CUDA tests of the DFlash draft block's non-causal windowed attention over a paged pool. + +``modeling_dflash.dflash_window_gather`` / ``dflash_windowed_block_attention`` serve a drafter +whose layers attend non-causally within a two-sided window (the DFlash 2 Beam drafter: a +4096-token window, eight-row blocks, 10 query / 2 KV heads of 128 per TP4 rank) on the TRTLLM +draft-attention backend, whose TRTLLM-Gen kernels have no such mask. Query ``j`` of a block +sits at slot-local position ``ctx + j`` and must see exactly the request's keys within +``window - 1`` positions on either side: its committed context and the current block, nothing +at or past ``ctx + block_size``, nothing of another request. + +What is checked, against independent integer and FP32 references (never the helper's own +mask): + +* every query's effective key-index set equals the integer rule at contexts on both sides of + the window, including the 4095/4096/4097 boundary and short contexts; +* boundary sentinels with one-hot values: a key at distance ``window - 1`` contributes, one at + distance ``window`` contributes nothing, so a missing or extra key changes the output by a + whole one-hot column rather than by an average's rounding; +* stale and foreign pages: high-scoring poisoned keys past the block and in other requests' + pages contribute nothing; a rewound block (smaller context after a rejection) excludes the + tail it no longer owns; +* the attention output against explicit FP32 masked softmax (reported, BF16 tolerance). + +Needs one CUDA device; no model files. +""" + +from __future__ import annotations + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_dflash import ( + dflash_window_gather, + dflash_windowed_block_attention, +) + +WINDOW = 4096 +BLOCK = 8 +PAGE = 32 +HQ, HKV, D = 10, 2, 128 +SCALE = D**-0.5 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _oracle(ctx: int, j: int) -> set: + """Key positions query ``j`` of a block at context ``ctx`` must see (integer rule).""" + q = ctx + j + return {k for k in range(0, ctx + BLOCK) if abs(q - k) <= WINDOW - 1} + + +def _pool(ctx_lens, *, poison: float = 0.0): + """A paged HND pool with each request on its own pages and an identity-free block table. + + Pages are dealt out interleaved across requests (request r owns pages r, r + R, ...), so a + request's pages are not contiguous and a wrong row of the table reads another request. + Returns ``(pool [pages, 2, HKV, PAGE, D], block_tables [R, pages_per_row])``. + """ + rows = len(ctx_lens) + pages_per_row = (max(ctx_lens) + BLOCK + PAGE - 1) // PAGE + 1 + total = rows * pages_per_row + pool = torch.full((total, 2, HKV, PAGE, D), poison, dtype=torch.bfloat16, device="cuda") + tables = torch.arange(total, dtype=torch.int32, device="cuda").view(pages_per_row, rows).t() + return pool, tables.contiguous() + + +def _write(pool, tables, row, position, k=None, v=None): + page = int(tables[row, position // PAGE]) + if k is not None: + pool[page, 0, :, position % PAGE] = k + if v is not None: + pool[page, 1, :, position % PAGE] = v + + +@pytest.mark.parametrize( + "ctx_lens", + [ + (1, 5, 52, 513), + (4087, 4088, 4089, 4095), + (4096, 4097, 4100, 5000), + (0, 4095, 4096, 4097), + ], + ids=["short", "below-window", "past-window", "boundary"], +) +def test_key_index_sets_equal_integer_rule(ctx_lens): + """Every query's gathered, unmasked key positions are exactly the integer rule's set.""" + _, tables = _pool(ctx_lens) + num_ctx = torch.tensor(ctx_lens, dtype=torch.int32, device="cuda") + gather = dflash_window_gather(tables, num_ctx, BLOCK, WINDOW, PAGE) + assert gather.positions.shape == (len(ctx_lens), WINDOW + BLOCK - 1) + assert gather.mask.shape == (len(ctx_lens), 1, BLOCK, WINDOW + BLOCK - 1) + positions = gather.positions.cpu() + mask = gather.mask[:, 0].cpu() + valid = gather.valid.cpu() + for row, ctx in enumerate(ctx_lens): + for j in range(BLOCK): + seen = set(positions[row][mask[row, j]].tolist()) + assert seen == _oracle(ctx, j), (ctx, j) + # Valid keys are the request's own: committed context and the current block. + assert set(positions[row][valid[row]].tolist()) == set(range(0, ctx + BLOCK)) & set( + positions[row].tolist() + ) + + +def test_invalid_positions_address_the_blocks_own_page(): + """A position outside the request's keys is redirected to its block's first slot.""" + ctx_lens = (3, 4200) + _, tables = _pool(ctx_lens) + num_ctx = torch.tensor(ctx_lens, dtype=torch.int32, device="cuda") + gather = dflash_window_gather(tables, num_ctx, BLOCK, WINDOW, PAGE) + for row, ctx in enumerate(ctx_lens): + own_page = int(tables[row, ctx // PAGE]) + invalid = ~gather.valid[row] + assert bool((gather.pages[row][invalid] == own_page).all()) + assert bool((gather.offsets[row][invalid] == ctx % PAGE).all()) + owned = set(tables[row].tolist()) + assert set(gather.pages[row].tolist()) <= owned + + +def _one_hot(column: int) -> torch.Tensor: + v = torch.zeros(HKV, D, dtype=torch.bfloat16, device="cuda") + v[:, column] = 1.0 + return v + + +def _sentinel_query(rows: int) -> torch.Tensor: + """Queries along e_0 scaled so that a key ``e_0`` scores 40 and a zero key scores 0.""" + q = torch.zeros(rows, BLOCK, HQ, D, dtype=torch.float32, device="cuda") + q[..., 0] = 40.0 / SCALE + return q.to(torch.bfloat16) + + +def _place_sentinels(pool, tables, row: int, ctx: int, *, stale: int = 0) -> None: + """Keys of request ``row`` at ``ctx - 4096 + m`` (m = 0..8) score 40 with value ``e_m``; + ``stale`` keys past its block (a rejected tail it no longer owns) score 80 with ``e_9``.""" + key = _one_hot(0) + for m in range(BLOCK + 1): + position = ctx - WINDOW + m + if position >= 0: + _write(pool, tables, row, position, k=key, v=_one_hot(m)) + for s in range(stale): + _write(pool, tables, row, ctx + BLOCK + s, k=2 * key, v=_one_hot(9)) + + +def _boundary_inputs(ctx: int): + pool, tables = _pool([ctx]) + _place_sentinels(pool, tables, 0, ctx) + return _sentinel_query(1), pool, tables + + +@pytest.mark.parametrize("ctx", [4095, 4096, 4097, 4103, 6000]) +def test_window_boundary_sentinels(ctx): + """Query j sees the sentinel keys at distance <= 4095 and none at distance 4096. + + The visible sentinels share the softmax almost evenly (each scores ~40 against ~4000 + zero-score keys), so column m of query j's output is ~1/(#visible) exactly for the + visible m and ~0 for the others; one extra or missing boundary key moves a whole column. + """ + q, pool, tables = _boundary_inputs(ctx) + num_ctx = torch.tensor([ctx], dtype=torch.int32, device="cuda") + gather = dflash_window_gather(tables, num_ctx, BLOCK, WINDOW, PAGE) + out = dflash_windowed_block_attention(q, pool, gather, SCALE).float()[0] + for j in range(BLOCK): + visible = [m for m in range(BLOCK + 1) if ctx - WINDOW + m >= 0 and m >= j + 1] + expected = torch.zeros(BLOCK + 1, device="cuda") + if visible: + expected[visible] = 1.0 / len(visible) + got = out[j, :, : BLOCK + 1] + assert torch.allclose(got, expected.expand_as(got), atol=2e-2), (ctx, j, got[0]) + + +@pytest.mark.parametrize("stale", [1, 8]) +def test_stale_tail_and_foreign_pages_contribute_nothing(stale): + """Poisoned keys past request 0's block (a rewound tail still in its pages) and every slot + of request 1's interleaved pages outscore all valid keys, so any leak would dominate: + request 0's output must carry none of their value columns.""" + ctx = 4100 + pool, tables = _pool([ctx + stale, ctx + stale]) + _place_sentinels(pool, tables, 0, ctx, stale=stale) + for page in tables[1].tolist(): + pool[page, 0, :, :, 0] = 2.0 + pool[page, 1] = _one_hot(10).unsqueeze(1).expand(HKV, PAGE, D) + num_ctx = torch.tensor([ctx, ctx], dtype=torch.int32, device="cuda") + gather = dflash_window_gather(tables, num_ctx, BLOCK, WINDOW, PAGE) + out = dflash_windowed_block_attention(_sentinel_query(2), pool, gather, SCALE).float() + assert float(out[0, :, :, 9].abs().max()) < 1e-3, "a key past the block leaked" + assert float(out[0, :, :, 10].abs().max()) < 1e-3, "another request's page leaked" + for j in range(BLOCK): + assert abs(float(out[0, j, 0, j + 1 : BLOCK + 1].sum()) - 1.0) < 2e-2, j + # Request 1's keys within its window are all poison: it sees only them. + assert float((out[1, :, :, 10] - 1.0).abs().max()) < 2e-2 + + +def _rewound_block(num_ctx: int): + """A block appended at 4092..4099, then a rejection that committed context 4095, plus a + far stale row at 4103..4107 (key scoring 80, value e_9); the next block is gathered at + context ``num_ctx``. Returns the block's output and gather.""" + pool, tables = _pool([4100]) + _place_sentinels(pool, tables, 0, 4095) + key = _one_hot(0) + for position in range(4103, 4108): + _write(pool, tables, 0, position, k=2 * key, v=_one_hot(9)) + gather = dflash_window_gather( + tables, torch.tensor([num_ctx], dtype=torch.int32, device="cuda"), BLOCK, WINDOW, PAGE + ) + out = dflash_windowed_block_attention(_sentinel_query(1), pool, gather, SCALE).float() + return out, gather + + +def _check_rewound(out, gather, committed: int) -> None: + """The block after the rejection reads no stale value and its keys end at its own rows.""" + assert float(out[0, :, :, 9].abs().max()) < 1e-3, "a stale row past the block was read" + assert int(gather.positions[0][gather.valid[0]].max()) == committed + BLOCK - 1, ( + "the gathered keys do not end at the block's own last row" + ) + + +def test_rewound_block_excludes_the_tail_it_no_longer_owns(): + """The next block at the committed context 4095 must neither see the superseded rows past + 4102 nor miss a key.""" + _check_rewound(*_rewound_block(4095), committed=4095) + + +def test_unrewound_draft_context_is_detected(): + """Fault control: a drafter that keeps the superseded block as committed context (context + 4100 where the rejection committed 4095) gathers the stale rows; the same check fails.""" + with pytest.raises(AssertionError, match="stale row"): + _check_rewound(*_rewound_block(4100), committed=4095) + + +def _fp32_reference(q, pool, tables, ctx_lens): + out = torch.empty(q.shape, dtype=torch.float32, device="cuda") + for row, ctx in enumerate(ctx_lens): + keys = list(range(0, ctx + BLOCK)) + pages = tables[row, torch.tensor(keys, device="cuda") // PAGE].long() + offsets = torch.tensor(keys, device="cuda") % PAGE + k = pool[pages, 0, :, offsets].float().repeat_interleave(HQ // HKV, dim=1) + v = pool[pages, 1, :, offsets].float().repeat_interleave(HQ // HKV, dim=1) + kpos = torch.tensor(keys, device="cuda") + for j in range(BLOCK): + allowed = (ctx + j - kpos).abs() <= WINDOW - 1 + s = torch.einsum("hd,khd->hk", q[row, j].float(), k) * SCALE + s = s.masked_fill(~allowed[None], float("-inf")) + out[row, j] = torch.einsum("hk,khd->hd", s.softmax(-1), v) + return out + + +def test_matches_fp32_masked_softmax(): + """Random data at four requests (the TP4 geometry): BF16 SDPA vs explicit FP32.""" + torch.manual_seed(0) + ctx_lens = (52, 513, 4095, 4097) + pool, tables = _pool(ctx_lens) + pool.normal_() + q = torch.randn(len(ctx_lens), BLOCK, HQ, D, dtype=torch.bfloat16, device="cuda") + num_ctx = torch.tensor(ctx_lens, dtype=torch.int32, device="cuda") + gather = dflash_window_gather(tables, num_ctx, BLOCK, WINDOW, PAGE) + out = dflash_windowed_block_attention(q, pool, gather, SCALE).float() + ref = _fp32_reference(q, pool, tables, ctx_lens) + rel = float((out - ref).abs().max() / ref.abs().max()) + assert rel < 1e-2, rel diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py index 5c196b997868..a563aab7e9ba 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py @@ -5539,6 +5539,7 @@ def _op_reference( causal: bool = True, e4m3: bool = False, layer: Optional[int] = None, + window: Optional[int] = None, nvfp4: bool = False, ) -> torch.Tensor: """`op.reference` over the history the env mirrored for these requests. @@ -5570,6 +5571,7 @@ def _op_reference( v_head_dim=d, softmax_scale=1.0 / math.sqrt(d), causal=causal, + window=window, e4m3_inputs=e4m3, e4m3_probabilities=e4m3, ) @@ -7113,10 +7115,11 @@ def _cell_paged(cell, arm) -> None: tokens_per_block=spec["tokens_per_block"], num_blocks=spec["num_blocks"], max_batch=spec["max_num_requests"], - max_blocks_per_seq=16, + max_blocks_per_seq=spec.get("max_blocks_per_seq", 16), pool_dtype=torch.float8_e4m3fn if fp8 else torch.bfloat16, quant_mode=spec["quant_mode"], kv_scaling_factor=spec["kv_scale"] if fp8 or nvfp4 else None, + attention_window_size=spec["window"], ) cached, new, nc = spec["cached"], spec["new"], spec["num_contexts"] p, paged = spec["predicted_tokens_per_seq"], spec["use_paged_context_fmha"] @@ -7154,7 +7157,10 @@ def _cell_paged(cell, arm) -> None: assert not (nvfp4 and not paged), "an NVFP4 cell runs the paged-context path" e4m3 = (fp8 or nvfp4) and (paged or nc == 0) causal = spec["mask_type"] == MASK_CAUSAL - _compare(out, _op_reference(env, qkv, list(new), rids, causal=causal, e4m3=e4m3, nvfp4=nvfp4)) + ref = _op_reference( + env, qkv, list(new), rids, causal=causal, e4m3=e4m3, window=spec["window"], nvfp4=nvfp4 + ) + _compare(out, ref) for rid in rids: env.check_cache(rid) env.check_unwritten_pool_zero() @@ -7296,3 +7302,150 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 _fp8_mla_r1_generation_decode( cell.spec["num_heads"], cell.spec.get("predicted_tokens_per_seq", 1) ) + + +# --------------------------------------------------------------------------------------------- +# Beam TP4 speculative verification at the 512-token sliding window +# --------------------------------------------------------------------------------------------- + +#: One TP4 rank of the Beam full model: 12 query / 2 KV heads of 128, 32-token pages. +BEAM_HQ, BEAM_HKV, BEAM_D, BEAM_WINDOW, BEAM_P = 12, 2, 128, 512, 8 + + +def _beam_indicator_qkv(positions: List[int]) -> torch.Tensor: + """The one-hot read-out at the Beam TP4 geometry: K = 0 (every unmasked logit is 0, so the + softmax is uniform over exactly the attended keys), V[t] = e_t for t < 128 and 0 after, so + output column d of a query row is the weight it gives key d.""" + q_w, kv_w = BEAM_HQ * BEAM_D, BEAM_HKV * BEAM_D + qkv = torch.zeros(len(positions), q_w + 2 * kv_w, dtype=torch.bfloat16, device="cuda") + qkv[:, :q_w] = torch.randn(len(positions), q_w, dtype=torch.bfloat16, device="cuda") + v = torch.zeros(len(positions), BEAM_HKV, BEAM_D, dtype=torch.bfloat16, device="cuda") + for i, t in enumerate(positions): + if t < BEAM_D: + v[i, :, t] = 1.0 + qkv[:, q_w + kv_w :] = v.reshape(len(positions), kv_w) + return qkv + + +def _assert_beam_row_weights(out, row: int, pos: int, window: int, label: str) -> None: + """Row ``row`` is the query at absolute position ``pos``: key d < 128 carries exactly + 1 / n inside ``[pos - window + 1, pos]`` (n = keys in that range) and exactly 0 outside.""" + w = out.view(-1, BEAM_HQ, BEAM_D)[row].float() + lo, n = max(0, pos - window + 1), min(pos + 1, window) + inside = [d for d in range(BEAM_D) if lo <= d <= pos] + outside = [d for d in range(BEAM_D) if not lo <= d <= pos] + if outside: + assert bool((w[:, outside] == 0).all()), f"{label}: a key outside [{lo}, {pos}] is attended" + expect = torch.full((BEAM_HQ,), 1.0 / n, dtype=torch.float32, device="cuda") + for d in inside: + torch.testing.assert_close( + w[:, d], + expect, + rtol=SWA_WEIGHT_RTOL, + atol=0.0, + msg=lambda m, d=d: f"{label}: key {d} weight off\n{m}", + ) + + +def _beam_verify_env(window: Optional[int]) -> _PagedAttnEnv: + return _PagedAttnEnv( + num_heads=BEAM_HQ, + num_kv_heads=BEAM_HKV, + head_dim=BEAM_D, + tokens_per_block=32, + num_blocks=64, + max_batch=1, + max_blocks_per_seq=24, + attention_window_size=window, + ) + + +def _beam_block_call(env: _PagedAttnEnv) -> torch.Tensor: + """One P = 8 verification block appended after the request's cached history.""" + start = env.cached_len(0) + return env.call_op( + _beam_indicator_qkv(list(range(start, start + BEAM_P))), + [BEAM_P], + 0, + [0], + MASK_CAUSAL, + ctx_lens_override=[env.prompt_lens[0]], + use_paged_context_fmha=True, + predicted_tokens_per_seq=BEAM_P, + ) + + +def _check_block(out: torch.Tensor, start: int, window: int, label: str) -> None: + """Every row of a block at positions start .. start + 7 reads out exactly its key set.""" + for r in range(BEAM_P): + _assert_beam_row_weights(out, r, start + r, window, f"{label} row {r}") + + +def _beam_prefill(env: _PagedAttnEnv, length: int) -> None: + env.add_request(0, length) + env.call_op( + _beam_indicator_qkv(list(range(length))), + [length], + 1, + [0], + MASK_CAUSAL, + ctx_lens_override=[length], + use_paged_context_fmha=True, + ) + + +def _rewind(env: _PagedAttnEnv, committed: int) -> None: + """Native rollback of the target cache: only ``committed`` positions stay history; the + rejected rows of the block are overwritten by the next block's append.""" + k, v = torch.cat(env.k_history[0]), torch.cat(env.v_history[0]) + env.k_history[0], env.v_history[0] = [k[:committed]], [v[:committed]] + + +@pytest.mark.parametrize("accepted", [2, 3, 4], ids=["commit511", "commit512", "commit513"]) +@pytest.mark.parametrize("window", [BEAM_WINDOW, None], ids=["sliding512", "full"]) +def test_beam_verify_block_straddles_the_window_and_rolls_back(window, accepted): + """A verification block at 508..515 over a 508-token history, ``accepted`` proposals + accepted (committed length 508 + 1 + accepted = 511/512/513), then the next block from the + committed length. On the sliding layer row r at position p keeps key p - 511 and drops key + p - 512 exactly; the full-attention companion keeps every key. After the rollback the pool + holds exactly the committed history and the new block, bitwise.""" + torch.manual_seed(512 + accepted) + env = _beam_verify_env(window) + effective = env.attention_window_size + _beam_prefill(env, 508) + _check_block(_beam_block_call(env), 508, effective, "first block") + committed = 508 + 1 + accepted + _rewind(env, committed) + _check_block(_beam_block_call(env), committed, effective, f"after committing {committed}") + env.check_cache(0) + + +def test_beam_verify_retained_rejected_tail_is_detected(): + """Fault control: a runtime that keeps a rejected block's rows as history (no rewind after + two accepted proposals) appends the next block at 516 instead of 511; the exact per-row + read-out the rollback test applies at the committed positions must fail.""" + torch.manual_seed(7) + env = _beam_verify_env(BEAM_WINDOW) + _beam_prefill(env, 508) + _check_block(_beam_block_call(env), 508, BEAM_WINDOW, "first block") + committed = 508 + 1 + 2 + out = _beam_block_call(env) # the fault: no _rewind(env, committed) + assert env.cached_len(0) == 508 + 2 * BEAM_P + with pytest.raises(AssertionError, match="weight off"): + _check_block(out, committed, BEAM_WINDOW, "retained tail") + + +def test_beam_verify_wrong_accepted_length_is_detected(): + """Fault control: an accounting that also counts the correction token's row as committed + (508 + 2 + 3 instead of 508 + 1 + 3 after three accepted proposals) appends the next block + one position late; the exact per-row read-out at the committed positions must fail.""" + torch.manual_seed(11) + env = _beam_verify_env(BEAM_WINDOW) + _beam_prefill(env, 508) + _check_block(_beam_block_call(env), 508, BEAM_WINDOW, "first block") + committed = 508 + 1 + 3 + _rewind(env, committed + 1) # the fault: one row more than the native contract commits + out = _beam_block_call(env) + assert env.cached_len(0) == committed + 1 + BEAM_P + with pytest.raises(AssertionError, match="weight off"): + _check_block(out, committed, BEAM_WINDOW, "wrong accepted length") diff --git a/tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/norm/test_modeling_v2_beam_rmsnorm_fp32.py b/tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/norm/test_modeling_v2_beam_rmsnorm_fp32.py new file mode 100644 index 000000000000..16a5acfdbbbb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/norm/test_modeling_v2_beam_rmsnorm_fp32.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_rmsnorm_fp32 entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_bf16__sm_103__tp4.ops.norm.beam_rmsnorm_fp32 import ( # noqa: E501 + beam_rmsnorm_fp32 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_rmsnorm_fp32 requires a CUDA device" + +SEEDS = (17, 41) + + +def _build(spec, seed): + """Hidden rows with a residual stream's spread (a few large channels), a weight near 1.""" + g = torch.Generator(device="cpu").manual_seed(seed) + x = torch.randn(spec["tokens"], spec["width"], generator=g) * 4.0 + x[:, :: spec["width"] // 8] *= 64.0 + w = 1.0 + 0.25 * torch.randn(spec["width"], generator=g) + return x.to(spec["dtype"]).cuda(), w.to(spec["dtype"]).cuda() + + +def _bf16_units_apart(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Distance between same-signed bf16 values in units of the last place.""" + assert a.dtype == b.dtype == torch.bfloat16 + assert bool(((a.float() >= 0) == (b.float() >= 0)).all()), "signs differ" + return (a.view(torch.int16).int() - b.view(torch.int16).int()).abs() + + +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:26] for c in op.CELLS]) +def test_certified_cells(cell, seed, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 unit of the fp64 reference: the fp32 steps carry a few fp32 units of + error, so the single rounding at the end may land on either neighbour of a boundary.""" + certifying(op) + x, w = _build(cell.spec, seed) + before = x.clone() + out = op(x, w, cell.spec["eps"]) + ref = op.reference(x, w, cell.spec["eps"]) + assert out.dtype == torch.bfloat16 and out.shape == x.shape + assert out.data_ptr() != x.data_ptr() and torch.equal(x, before) + assert int(_bf16_units_apart(out, ref).max()) <= 1 + + +def test_rounds_once_at_the_end() -> None: + """Rounding the normalized row to bf16 before the weight is a different answer: it + rounds twice, and on these rows some outputs move by a bf16 unit.""" + spec = op.CELLS[2].spec + x, w = _build(spec, seed=17) + out = op(x, w, spec["eps"]) + x32 = x.float() + rms = x32.pow(2).mean(dim=-1, keepdim=True).add(spec["eps"]).rsqrt() + rounded_twice = (x32 * rms).to(torch.bfloat16) * w + assert torch.equal(out, (x32 * rms * w).to(torch.bfloat16)) + assert bool((out != rounded_twice).any()) + + +@pytest.mark.parametrize( + "case,match", + [ + ("fp32_weight", "must be torch.bfloat16"), + ("narrow_weight", "weight must be"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + x, w = _build(op.CELLS[1].spec, seed=17) + w = w.float() if case == "fp32_weight" else w[:1] + with pytest.raises(AssertionError, match=match): + with validating(op): + op(x, w, 1e-6) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/test_checkpoint_layout.py b/tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/test_checkpoint_layout.py new file mode 100644 index 000000000000..18922e2c4e86 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/test_checkpoint_layout.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Generated-data tests for the Beam checkpoint layout component: bf16 / sm_103 / tp4. + +No model files. The layout tests use the architecture's configuration values only (no +tensors). The loader tests write small safetensors directories with this file's own writer +and compare every loaded byte with an expected value derived independently here: E4M3 +decoded from its bit fields, the product with the BF16 row scale formed exactly, and the +BF16 rounding done on the float32 bit pattern (round to nearest, ties to even). Every rank +of the four-rank layout is exercised, with tiles small enough to cross row seams. + +Four-rank execution (collectives, real checkpoint loading, KV cache) needs four GPUs and a +checkpoint and is not part of this file. + +Invocation:: + + env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE \\ + python -m pytest tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4 +""" + +from __future__ import annotations + +import json +import math +import os +import struct + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_bf16__sm_103__tp4 import ( # noqa: E501 + layout, + loader, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_bf16__sm_103__tp4.layout import ( # noqa: E501 + BF16, + F32, + FP8, + TensorSpec, +) + +NUM_LAYERS = 52 +FULL = {5, 11, 17, 23, 29, 35, 41, 47, 51} + + +def model_config() -> dict: + """The architecture's configuration, as its config.json states it (values only).""" + cfg = dict(layout.SUPPORTED_CONFIG) + cfg["layer_types"] = [ + "full_attention" if i in FULL else "sliding_attention" for i in range(NUM_LAYERS) + ] + cfg["mlp_layer_types"] = ["dense"] * 3 + ["sparse"] * (NUM_LAYERS - 3) + cfg["rope_parameters"] = { + "full_attention": None, + "sliding_attention": {"rope_theta": 10000.0, "rope_type": "default"}, + } + return cfg + + +# ------------------------------------------------------------------ independent expected values + + +def e4m3(code: int) -> float: + """E4M3 (fn) value of one byte from its bit fields: bias 7, subnormals, no infinities.""" + sign = -1.0 if code & 0x80 else 1.0 + exp, man = (code >> 3) & 0xF, code & 0x7 + if exp == 0xF and man == 0x7: + return math.nan + if exp == 0: + return sign * man / 8.0 * 2.0**-6 + return sign * (1.0 + man / 8.0) * 2.0 ** (exp - 7) + + +def bf16_bits_of_float(x: float) -> int: + """BF16 bits of the float32 value nearest ``x`` (exact here), rounded to nearest even.""" + (u,) = struct.unpack("> 16) & 1)) >> 16) & 0xFFFF + + +def bf16_value(bits: int) -> float: + return struct.unpack(" None: + """``shards``: file name -> list of (name, dtype, shape, raw bytes). Writes index + shards.""" + os.makedirs(path, exist_ok=True) + weight_map = {} + for filename, entries in shards.items(): + header, payload, offset = {}, b"", 0 + for name, dtype, shape, raw in entries: + header[name] = { + "dtype": dtype, + "shape": list(shape), + "data_offsets": [offset, offset + len(raw)], + } + payload += raw + offset += len(raw) + weight_map[name] = filename + blob = json.dumps(header).encode() + blob += b" " * (-len(blob) % 8) + with open(os.path.join(path, filename), "wb") as f: + f.write(struct.pack(" list: + """Codes cycling through every finite E4M3 byte (signs, zeros, subnormals, 448).""" + return [ + [FP8_CODES[(start + r * cols + c) % len(FP8_CODES)] for c in range(cols)] + for r in range(rows) + ] + + +def row_scales(rows: int, start: int) -> list: + """Distinct, non-uniform BF16 per-row scales (bits), as a quantizer's amax/448 would be.""" + return [ + bf16_bits_of_float(2.0 ** (-(r % 5) - 3) * (1.0 + (start + 3 * r) % 7 / 8.0)) + for r in range(rows) + ] + + +def fp8_entries(name: str, codes: list, scales: list) -> list: + rows, cols = len(codes), len(codes[0]) + raw = bytes(c for row in codes for c in row) + sraw = b"".join(struct.pack(" list: + """BF16 bits of bf16(fp32(e4m3(q)) * fp32(scale[row])), independently of torch.""" + return [ + [bf16_bits_of_float(e4m3(q) * bf16_value(scales[r])) for q in row] + for r, row in enumerate(codes) + ] + + +# A generated layout: row-split and column-split FP8, replicated BF16, row-split F32, on two shards. +SPECS = { + "a.rows.weight": TensorSpec("a.rows.weight", (8, 12), FP8, "gen", 0, None, 0), + "b.cols.weight": TensorSpec("b.cols.weight", (12, 16), FP8, "gen", 0, None, 1), + "c.norm.weight": TensorSpec("c.norm.weight", (10,), BF16, "gen", 0, None, None), + "d.router.weight": TensorSpec("d.router.weight", (8, 4), F32, "gen", 0, None, 0), +} +A_CODES, A_SCALES = fp8_matrix(8, 12, 0), row_scales(8, 0) +B_CODES, B_SCALES = fp8_matrix(12, 16, 97), row_scales(12, 5) +C_BITS = [ + bf16_bits_of_float(v) + for v in (-0.0, 1.5, -2.25, 3.0e-40, 448.0, 1.0 / 3, 7.0, -1e5, 0.0, 65504.0) +] +D_VALUES = [((-1) ** i) * (i + 0.5) / 3.0 for i in range(32)] + + +def generated_shards() -> dict: + return { + "model-00001-of-00002.safetensors": fp8_entries("a.rows.weight", A_CODES, A_SCALES) + + [("c.norm.weight", BF16, (10,), b"".join(struct.pack(" str: + path = str(tmp_path / "ckpt") + write_checkpoint(path, generated_shards()) + return path + + +def _bits16(t: torch.Tensor) -> list: + return t.contiguous().view(torch.int16).cpu().to(torch.int32).bitwise_and(0xFFFF).tolist() + + +# ------------------------------------------------------------------ layout + + +def test_layout_counts_intervals_and_banks() -> None: + specs = layout.tensor_specs(model_config()) + assert len(specs) == 3 + 3 * 15 + 49 * 782 + fp8 = [s for s in specs.values() if s.stored_dtype == FP8] + assert len(fp8) == 52 * 5 + 3 * 3 + 49 * 768 + layout.check_partition(specs) + want = { + "model.embed_tokens.weight": (0, 37984), + "lm_head.weight": (0, 37984), + "model.layers.0.self_attn.q_proj.weight": (0, 1536), + "model.layers.0.self_attn.gate_proj.weight": (0, 1536), + "model.layers.0.self_attn.k_proj.weight": (0, 256), + "model.layers.0.self_attn.v_proj.weight": (0, 256), + "model.layers.0.self_attn.o_proj.weight": (1, 1536), + "model.layers.0.self_attn.gate_bias": (0, 12), + "model.layers.0.mlp.gate_proj.weight": (0, 5120), + "model.layers.0.mlp.down_proj.weight": (1, 5120), + "model.layers.51.mlp.experts.255.up_proj.weight": (0, 640), + "model.layers.51.mlp.experts.255.down_proj.weight": (1, 640), + } + for name, (dim, size) in want.items(): + spec = specs[name] + got = [spec.interval(r) for r in range(4)] + assert got == [(dim, r * size, (r + 1) * size) for r in range(4)], name + for name in ( + "model.norm.weight", + "model.layers.3.mlp.gate.weight", + "model.layers.3.self_attn.q_norm.weight", + ): + assert all(specs[name].interval(r) is None for r in range(4)) + router = specs["model.layers.3.mlp.gate.e_score_correction_bias"] + assert router.stored_dtype == F32 and router.effective_dtype == F32 + expert = specs["model.layers.7.mlp.experts.200.down_proj.weight"] + assert expert.bank == "layers.7.experts.down" and expert.expert == 200 + assert expert.local_shape(2) == (5120, 640) and expert.scale_name.endswith( + "down_proj.weight_scale" + ) + assert ( + specs["lm_head.weight"].stored_dtype == BF16 and specs["lm_head.weight"].scale_name is None + ) + + +def test_gqa_pairing_and_ledger_rows() -> None: + specs = layout.tensor_specs(model_config(), layers=(5,)) + q, k = ( + specs["model.layers.5.self_attn.q_proj.weight"], + specs["model.layers.5.self_attn.k_proj.weight"], + ) + for rank in range(4): + _, q0, q1 = q.interval(rank) + _, k0, k1 = k.interval(rank) + q_heads = range(q0 // 128, q1 // 128) + kv_heads = set(range(k0 // 128, k1 // 128)) + assert {h // 6 for h in q_heads} == kv_heads == {2 * rank, 2 * rank + 1} + rows = {row["name"]: row for row in layout.ownership_ledger(specs)} + assert len(rows) == len(specs) + o = rows["model.layers.5.self_attn.o_proj.weight"]["ranks"] + assert [(r["dim"], r["start"], r["stop"]) for r in o] == [ + (1, 1536 * i, 1536 * (i + 1)) for i in range(4) + ] + bank = rows["model.layers.5.mlp.experts.9.gate_proj.weight"]["ranks"][3] + assert (bank["bank"], bank["slot"], bank["start"], bank["stop"]) == ( + "layers.5.experts.gate", + 9, + 1920, + 2560, + ) + assert all(r["replicated"] for r in rows["model.layers.5.mlp.gate.weight"]["ranks"]) + + +@pytest.mark.parametrize( + "change", + [ + {"num_experts": 255}, + {"routed_scaling_factor": 2.0}, + {"tie_word_embeddings": True}, + {"mlp_layer_types": ["sparse"] + ["dense"] * 2 + ["sparse"] * 49}, + {"layer_types": ["sliding_attention"] * 51}, + { + "rope_parameters": { + "full_attention": None, + "sliding_attention": {"rope_theta": 1e6, "rope_type": "default"}, + } + }, + {"num_key_value_heads": 4}, + ], +) +def test_inconsistent_configuration_is_rejected(change) -> None: + cfg = model_config() + cfg.update(change) + with pytest.raises(ValueError): + layout.tensor_specs(cfg) + + +@pytest.mark.parametrize("rank, world", [(4, 4), (-1, 4), (0, 2), (0, 8), (True, 4), (1.0, 4)]) +def test_invalid_rank_or_world_is_rejected(rank, world) -> None: + with pytest.raises(ValueError): + layout.check_rank(rank, world) + spec = layout.tensor_specs(model_config(), layers=(0,))[ + "model.layers.0.self_attn.q_proj.weight" + ] + with pytest.raises(ValueError): + spec.interval(rank, world) + + +def test_unknown_layers_are_rejected() -> None: + for layers in ((52,), (-1,), (3, 3), ("3",)): + with pytest.raises(ValueError): + layout.tensor_specs(model_config(), layers=layers) + + +# ------------------------------------------------------------------ catalog validation + + +def test_catalog_accepts_the_exact_checkpoint(checkpoint) -> None: + catalog = loader.SafetensorsCatalog(checkpoint) + catalog.validate(SPECS, complete=True) + assert set(catalog.entries) == set(SPECS) | {"a.rows.weight_scale", "b.cols.weight_scale"} + + +def _mutated(entries_by_file: dict, how: str) -> dict: + shards = {f: list(e) for f, e in entries_by_file.items()} + first, second = list(shards) + if how == "missing_scale": + shards[first] = [e for e in shards[first] if e[0] != "a.rows.weight_scale"] + elif how == "extra_tensor": + shards[second].append(("e.unsupported.weight", BF16, (2,), b"\0" * 4)) + elif how == "scale_vector": + shards[first] = [ + (n, d, (8,) if n == "a.rows.weight_scale" else s, raw) for n, d, s, raw in shards[first] + ] + elif how == "wrong_dtype": + shards[second] = [ + (n, BF16 if n == "d.router.weight" else d, (8, 8) if n == "d.router.weight" else s, raw) + for n, d, s, raw in shards[second] + ] + elif how == "wrong_shape": + shards[first] = [ + (n, d, (12, 8) if n == "a.rows.weight" else s, raw) for n, d, s, raw in shards[first] + ] + elif how == "transposed_scale": + shards[second] = [ + (n, d, (1, 12) if n.endswith("_scale") else s, raw) for n, d, s, raw in shards[second] + ] + return shards + + +@pytest.mark.parametrize( + "how", + [ + "missing_scale", + "extra_tensor", + "scale_vector", + "wrong_dtype", + "wrong_shape", + "transposed_scale", + ], +) +def test_catalog_rejects_malformed_metadata(tmp_path, how) -> None: + path = str(tmp_path / how) + write_checkpoint(path, _mutated(generated_shards(), how)) + with pytest.raises(loader.CheckpointFormatError): + loader.SafetensorsCatalog(path).validate(SPECS, complete=True) + + +def test_catalog_rejects_index_header_disagreement_and_truncation(checkpoint) -> None: + index_path = os.path.join(checkpoint, loader.INDEX_FILE) + index = json.load(open(index_path)) + index["weight_map"]["c.norm.weight"] = "model-00002-of-00002.safetensors" + json.dump(index, open(index_path, "w")) + with pytest.raises(loader.CheckpointFormatError): + loader.SafetensorsCatalog(checkpoint) + write_checkpoint(checkpoint, generated_shards()) + shard = os.path.join(checkpoint, "model-00002-of-00002.safetensors") + with open(shard, "rb+") as f: + f.truncate(os.path.getsize(shard) - 3) + with pytest.raises(loader.CheckpointFormatError): + loader.SafetensorsCatalog(checkpoint) + + +# ------------------------------------------------------------------ streaming reads on CUDA + +cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +@cuda +@pytest.mark.parametrize("tile_bytes", [16, 40, loader.TILE_BYTES]) +def test_fp8_slices_match_independent_bytes_on_every_rank(checkpoint, tile_bytes) -> None: + catalog = loader.SafetensorsCatalog(checkpoint) + catalog.validate(SPECS) + reader = loader.TileReader(tile_bytes=tile_bytes, pin=True) + want_a, want_b = expected_fp8(A_CODES, A_SCALES), expected_fp8(B_CODES, B_SCALES) + try: + for rank in range(4): + a = torch.empty( + SPECS["a.rows.weight"].local_shape(rank), dtype=torch.bfloat16, device="cuda" + ) + loader.read_slice_into(catalog, reader, SPECS["a.rows.weight"], rank, a) + assert _bits16(a) == want_a[2 * rank : 2 * rank + 2] + b = torch.empty( + SPECS["b.cols.weight"].local_shape(rank), dtype=torch.bfloat16, device="cuda" + ) + loader.read_slice_into(catalog, reader, SPECS["b.cols.weight"], rank, b) + assert _bits16(b) == [row[4 * rank : 4 * rank + 4] for row in want_b] + finally: + reader.close() + + +@cuda +def test_fp8_rule_is_a_row_scale_not_a_cast_or_transpose(checkpoint) -> None: + """The fixtures distinguish the contract from a plain cast, a reciprocal and a column scale.""" + want = expected_fp8(B_CODES, B_SCALES) + plain = [[bf16_bits_of_float(e4m3(q)) for q in row] for row in B_CODES] + recip = [ + [bf16_bits_of_float(e4m3(q) / bf16_value(B_SCALES[r])) for q in row] + for r, row in enumerate(B_CODES) + ] + assert want != plain and want != recip + catalog = loader.SafetensorsCatalog(checkpoint) + reader = loader.TileReader(tile_bytes=64) + try: + b = torch.empty(SPECS["b.cols.weight"].local_shape(3), dtype=torch.bfloat16, device="cuda") + loader.read_slice_into(catalog, reader, SPECS["b.cols.weight"], 3, b) + finally: + reader.close() + assert _bits16(b) == [row[12:16] for row in want] + + +@cuda +def test_bf16_and_f32_tensors_copy_byte_for_byte(checkpoint) -> None: + catalog = loader.SafetensorsCatalog(checkpoint) + reader = loader.TileReader(tile_bytes=24) + try: + for rank in range(4): + c = torch.empty(10, dtype=torch.bfloat16, device="cuda") + loader.read_slice_into(catalog, reader, SPECS["c.norm.weight"], rank, c) + assert _bits16(c) == C_BITS + d = torch.empty(2, 4, dtype=torch.float32, device="cuda") + loader.read_slice_into(catalog, reader, SPECS["d.router.weight"], rank, d) + want = torch.tensor(D_VALUES, dtype=torch.float32).view(8, 4)[2 * rank : 2 * rank + 2] + assert torch.equal(d.cpu().view(torch.int32), want.view(torch.int32)) + finally: + reader.close() + + +@cuda +def test_wrong_destination_and_short_read_raise(checkpoint) -> None: + catalog = loader.SafetensorsCatalog(checkpoint) + reader = loader.TileReader(tile_bytes=32) + try: + wrong = torch.empty(2, 12, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError): + loader.read_slice_into(catalog, reader, SPECS["a.rows.weight"], 0, wrong) + shard = os.path.join(checkpoint, "model-00002-of-00002.safetensors") + with open(shard, "rb+") as f: + f.truncate(os.path.getsize(shard) - 64) + d = torch.empty(2, 4, dtype=torch.float32, device="cuda") + with pytest.raises(OSError): + loader.read_slice_into(catalog, reader, SPECS["d.router.weight"], 3, d) + finally: + reader.close() + + +# ------------------------------------------------------------------ component guards + + +def _sm103() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 3) + + +@pytest.mark.skipif(not _sm103(), reason="the component is built for sm_103") +def test_mapping_rejections_and_weight_state_guard() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_bf16__sm_103__tp4.model import ( # noqa: E501 + BeamCheckpointLayoutBf16Sm103Tp4, + ) + from tensorrt_llm.mapping import Mapping + + cfg = model_config() + for bad in ( + dict(world_size=2, rank=0, tp_size=2, moe_tp_size=2), + dict(world_size=4, rank=0, tp_size=4, moe_tp_size=1, moe_ep_size=4), + ): + with pytest.raises(ValueError): + BeamCheckpointLayoutBf16Sm103Tp4( + cfg, Mapping(gpus_per_node=4, pp_size=1, **bad), layers=(0,) + ) + mapping = Mapping( + world_size=4, rank=1, gpus_per_node=4, tp_size=4, pp_size=1, moe_tp_size=4, moe_ep_size=1 + ) + model = BeamCheckpointLayoutBf16Sm103Tp4(cfg, mapping, layers=(0,)) + assert model.vocab_interval == (37984, 75968) + assert tuple(model.tensor("model.layers.0.self_attn.k_proj.weight").shape) == (256, 5120) + state = model.weight_state() + with pytest.raises(RuntimeError): # no complete load: unusable + model.check_weight_state(state) + with pytest.raises(RuntimeError): + model.embed(torch.zeros(1, dtype=torch.int64, device="cuda")) + # White-box: mark usable to exercise the guard without a four-rank load. + model.usable = True + model.check_weight_state(state) + model.tensor("model.layers.0.input_layernorm.weight").add_(0) # in place: version changes + with pytest.raises(RuntimeError): + model.check_weight_state(state) + state = model.weight_state() + name = "model__layers__0__self_attn__q_norm__weight" + old = getattr(model, name) + setattr( + model, name, torch.nn.Parameter(old.detach().clone(), requires_grad=False) + ) # same-shape replacement + with pytest.raises(RuntimeError): + model.check_weight_state(state) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/test_checkpoint_layout_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/test_checkpoint_layout_nvfp4.py new file mode 100644 index 000000000000..9803cd5484ec --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/test_checkpoint_layout_nvfp4.py @@ -0,0 +1,624 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Generated-data tests for the Beam NVFP4-expert checkpoint layout: nvfp4 / sm_103 / tp4. + +No model files. The layout and recipe tests use the architecture's configuration values only. +The loader tests write small safetensors directories with this file's own writer and compare +every loaded byte with an independent slice of the generated bytes, on every rank of the +four-rank layout, with tiles small enough to cross row seams. The CUDA tests compare the +production block-scale relayout with this file's own 128x4 offset formula, and the NVFP4 expert +provider with this file's own decoder (E2M1 table, E4M3 bit fields, low nibble = even K index), +activation quantizer and FP64-accumulation GEMM on an asymmetric generated expert. + +Four-rank execution (collectives, real checkpoint loading, KV cache) needs four GPUs and a +checkpoint and is not part of this file. + +Invocation:: + + env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE \\ + python -m pytest tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4 +""" + +from __future__ import annotations + +import json +import math +import os +import struct + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_bf16__sm_103__tp4 import ( # noqa: E501 + layout as bf16_layout, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_nvfp4__sm_103__tp4 import ( # noqa: E501 + layout, + loader, +) + +NUM_LAYERS = 52 +FULL = {5, 11, 17, 23, 29, 35, 41, 47, 51} +SPARSE = range(3, NUM_LAYERS) + + +def model_config() -> dict: + """The architecture's configuration, as its config.json states it (values only).""" + cfg = dict(bf16_layout.SUPPORTED_CONFIG) + cfg["layer_types"] = [ + "full_attention" if i in FULL else "sliding_attention" for i in range(NUM_LAYERS) + ] + cfg["mlp_layer_types"] = ["dense"] * 3 + ["sparse"] * (NUM_LAYERS - 3) + cfg["rope_parameters"] = { + "full_attention": None, + "sliding_attention": {"rope_theta": 10000.0, "rope_type": "default"}, + } + return cfg + + +def recipe() -> dict: + """A ModelOpt-style NVFP4 recipe excluding every higher-precision Linear module.""" + patterns = ["lm_head", "model.embed_tokens", "model.layers.0*", "model.layers.1.*"] + patterns.append("model.layers.2.*") + for layer in SPARSE: + patterns += [f"model.layers.{layer}.mlp.gate", f"model.layers.{layer}.self_attn*"] + return { + "producer": {"name": "modelopt", "version": "generated"}, + "quantization": { + "quant_algo": "NVFP4", + "kv_cache_quant_algo": None, + "group_size": 16, + "exclude_modules": patterns, + }, + } + + +# ------------------------------------------------------------------ independent format helpers + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def e2m1(code: int) -> float: + """E2M1 value of one 4-bit code: bit 3 sign, bits 0-2 index the magnitude table.""" + return (-1.0 if code & 8 else 1.0) * _E2M1[code & 7] + + +def e4m3(code: int) -> float: + """E4M3 (fn) value of one byte from its bit fields: bias 7, subnormals, no infinities.""" + sign = -1.0 if code & 0x80 else 1.0 + exp, man = (code >> 3) & 0xF, code & 0x7 + if exp == 0xF and man == 0x7: + return math.nan + if exp == 0: + return sign * man / 8.0 * 2.0**-6 + return sign * (1.0 + man / 8.0) * 2.0 ** (exp - 7) + + +_E4M3_TABLE = torch.tensor([e4m3(b) for b in range(256)], dtype=torch.float64) +_E2M1_TABLE = torch.tensor([e2m1(c) for c in range(16)], dtype=torch.float64) + + +def unpack(packed: torch.Tensor) -> torch.Tensor: + """``[R, K/2]`` bytes -> ``[R, K]`` codes; the low nibble is the even K index.""" + p = packed.long() + return torch.stack((p & 0xF, p >> 4), dim=-1).reshape(p.shape[0], -1) + + +def decode(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """Packed codes x block scales, float64 ``[R, K]`` (global scale not applied).""" + codes = unpack(packed.cpu()) + vals = _E2M1_TABLE[codes] + sf = _E4M3_TABLE[scales.cpu().long()] + return (vals.view(vals.shape[0], -1, 16) * sf.unsqueeze(-1)).view(vals.shape) + + +def swizzle_offset(m: int, k: int, cols: int) -> int: + """Byte offset of block scale (row m, column k) in the 128x4 swizzled layout.""" + k_tiles = -(-cols // 4) + return ( + (m // 128) * k_tiles * 512 + (k // 4) * 512 + (m % 32) * 16 + ((m % 128) // 32) * 4 + k % 4 + ) + + +def unswizzle(flat: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + idx = torch.tensor( + [swizzle_offset(m, k, cols) for m in range(rows) for k in range(cols)], dtype=torch.long + ) + return flat.reshape(-1).cpu()[idx].view(rows, cols) + + +# ------------------------------------------------------------------ layout and recipe + + +def test_layout_counts_dtypes_and_rank_intervals() -> None: + specs = layout.tensor_specs(model_config()) + quant = [s for s in specs.values() if isinstance(s, layout.QuantSpec)] + assert len(quant) == 49 * 256 * 3 + entries = layout.stored_entries(specs) + by_dtype = {} + for dtype, _ in entries.values(): + by_dtype[dtype] = by_dtype.get(dtype, 0) + 1 + # 636 higher-precision BF16 tensors, 98 FP32 router tensors, one quartet per projection + assert by_dtype == { + "BF16": 636, + "F32": 98 + 2 * len(quant), + "U8": len(quant), + "F8_E4M3": len(quant), + } + for name, spec in specs.items(): + if isinstance(spec, layout.QuantSpec): + continue + if ".mlp.gate." in name: + assert spec.stored_dtype == "F32" + else: + assert spec.stored_dtype == "BF16", name + gate = specs["model.layers.3.mlp.experts.7.gate_proj"] + down = specs["model.layers.51.mlp.experts.255.down_proj"] + stored = gate.stored() + assert stored["model.layers.3.mlp.experts.7.gate_proj.weight"] == ("U8", (2560, 2560)) + assert stored["model.layers.3.mlp.experts.7.gate_proj.weight_scale"] == ("F8_E4M3", (2560, 320)) + assert stored["model.layers.3.mlp.experts.7.gate_proj.weight_scale_2"] == ("F32", ()) + assert stored["model.layers.3.mlp.experts.7.gate_proj.input_scale"] == ("F32", ()) + for r in range(4): + assert gate.interval(r) == (0, 640 * r, 640 * (r + 1)) + assert gate.packed_interval(r) == (0, 640 * r, 640 * (r + 1)) + assert gate.scale_interval(r) == (0, 640 * r, 640 * (r + 1)) + assert gate.local_packed_shape(r) == (640, 2560) + assert gate.local_scale_shape(r) == (640, 320) + assert down.interval(r) == (1, 640 * r, 640 * (r + 1)) + assert down.packed_interval(r) == (1, 320 * r, 320 * (r + 1)) + assert down.scale_interval(r) == (1, 40 * r, 40 * (r + 1)) + assert down.local_packed_shape(r) == (5120, 320) + assert down.local_scale_shape(r) == (5120, 40) + assert gate.bank == "layers.3.experts.gate" and down.bank == "layers.51.experts.down" + layout.check_partition(specs) + assert layout.swizzled_numel(640, 320) == 640 * 320 + assert layout.swizzled_numel(17, 40) == 128 * 40 + + +def test_rank_slices_must_fall_on_whole_blocks() -> None: + # 2600 columns split four ways is 650, which is not a whole number of 16-element blocks. + spec = layout.QuantSpec("x.down_proj", (5120, 2600), 3, 0, "down_proj", 1) + with pytest.raises(ValueError, match="whole 16-element blocks"): + spec.interval(0) + + +def test_ledger_rows_record_packed_and_scale_ownership() -> None: + specs = layout.tensor_specs(model_config(), layers=(5,)) + rows = {row["name"]: row for row in layout.ownership_ledger(specs)} + row = rows["model.layers.5.mlp.experts.200.down_proj"] + assert [r["packed"] for r in row["ranks"]] == [[320 * r, 320 * (r + 1)] for r in range(4)] + assert [r["scale"] for r in row["ranks"]] == [[40 * r, 40 * (r + 1)] for r in range(4)] + assert all(r["slot"] == 200 and r["bank"] == "layers.5.experts.down" for r in row["ranks"]) + assert all(r["scalars"] == "replicated" for r in row["ranks"]) + assert rows["model.layers.5.self_attn.q_proj.weight"]["stored_dtype"] == "BF16" + + +def test_recipe_is_reconciled_against_module_names() -> None: + specs = layout.tensor_specs(model_config()) + summary = layout.check_recipe(recipe(), specs) + assert summary["quantized_modules"] == 49 * 256 * 3 + # 52 x 5 attention projections, 3 x 3 dense projections, 49 routers, the head + assert summary["excluded_linear_modules"] == 52 * 5 + 9 + 49 + 1 + + +@pytest.mark.parametrize( + "mutate, match", + [ + (lambda q: q["quantization"].update(quant_algo="FP8"), "quant_algo"), + (lambda q: q["quantization"].update(group_size=32), "group_size"), + (lambda q: q["quantization"].update(kv_cache_quant_algo="FP8"), "kv_cache_quant_algo"), + ( + lambda q: q["quantization"]["exclude_modules"].append("model.layers.9.mlp.experts.3*"), + "excludes quantized experts", + ), + (lambda q: q["quantization"]["exclude_modules"].remove("lm_head"), "not excluded"), + ( + lambda q: q["quantization"]["exclude_modules"].remove("model.layers.7.mlp.gate"), + "not excluded", + ), + (lambda q: q["quantization"].update(exclude_modules="lm_head"), "list of strings"), + (lambda q: q.pop("quantization"), "quantization"), + ], +) +def test_recipe_rejections(mutate, match) -> None: + q = recipe() + mutate(q) + with pytest.raises(ValueError, match=match): + layout.check_recipe(q, layout.tensor_specs(model_config(), layers=(0, 7, 9))) + + +# ------------------------------------------------------------------ generated checkpoints + +_SIZE = {"U8": 1, "F8_E4M3": 1, "BF16": 2, "F32": 4, "I8": 1} + + +def write_checkpoint(path: str, shards: dict) -> None: + """``shards``: file -> {name: (dtype, shape, raw bytes)}; writes the shards and the index.""" + os.makedirs(path, exist_ok=True) + weight_map = {} + for filename, tensors in shards.items(): + header, payload, offset = {}, [], 0 + for name, (dtype, shape, raw) in tensors.items(): + header[name] = { + "dtype": dtype, + "shape": list(shape), + "data_offsets": [offset, offset + len(raw)], + } + payload.append(raw) + offset += len(raw) + weight_map[name] = filename + blob = json.dumps(header).encode() + with open(os.path.join(path, filename), "wb") as f: + f.write(struct.pack(" torch.Tensor: + g = torch.Generator().manual_seed(seed) + return torch.randint(0, 256, (rows, cols), dtype=torch.uint8, generator=g) + + +@pytest.fixture +def quartets(tmp_path) -> tuple: + """One expert's gate and down quartets (real geometry), split across two shards.""" + p = "model.layers.3.mlp.experts.7." + gate_w, gate_s = _bytes(2560, 2560, 1), _bytes(2560, 320, 2) & 0x7E + down_w, down_s = _bytes(5120, 1280, 3), _bytes(5120, 160, 4) & 0x7E + scalars = { + p + "gate_proj.weight_scale_2": 3.0517578125e-05, + p + "gate_proj.input_scale": 0.0123, + p + "down_proj.weight_scale_2": 1.5, + p + "down_proj.input_scale": 2.0**-7, + } + shards = { + "a.safetensors": { + p + "gate_proj.weight": ("U8", (2560, 2560), gate_w.numpy().tobytes()), + p + "gate_proj.weight_scale": ("F8_E4M3", (2560, 320), gate_s.numpy().tobytes()), + **{n: ("F32", (), struct.pack(" None: + path, expected, t, _ = quartets + catalog = loader.SafetensorsCatalog(path) + loader.validate(catalog, expected) + spec_g = layout.QuantSpec( + "model.layers.3.mlp.experts.7.gate_proj", (2560, 5120), 3, 7, "gate_proj", 0 + ) + spec_d = layout.QuantSpec( + "model.layers.3.mlp.experts.7.down_proj", (5120, 2560), 3, 7, "down_proj", 1 + ) + reader = loader.TileReader(tile_bytes, pin=False) + try: + for r in range(4): + _, a, b = spec_g.packed_interval(r) + dst = torch.empty(spec_g.local_packed_shape(r), dtype=torch.uint8) + loader.read_packed_into(catalog, reader, spec_g.name + ".weight", (a, b), None, dst) + assert torch.equal(dst, t["gate_w"][640 * r : 640 * (r + 1)]) + dst = torch.empty(spec_g.local_scale_shape(r), dtype=torch.uint8) + loader.read_packed_into( + catalog, reader, spec_g.name + ".weight_scale", (a, b), None, dst + ) + assert torch.equal(dst, t["gate_s"][640 * r : 640 * (r + 1)]) + _, a, b = spec_d.packed_interval(r) + dst = torch.empty(spec_d.local_packed_shape(r), dtype=torch.uint8) + loader.read_packed_into( + catalog, reader, spec_d.name + ".weight", (0, 5120), (a, b), dst + ) + assert torch.equal(dst, t["down_w"][:, 320 * r : 320 * (r + 1)]) + _, a, b = spec_d.scale_interval(r) + dst = torch.empty(spec_d.local_scale_shape(r), dtype=torch.uint8) + loader.read_packed_into( + catalog, reader, spec_d.name + ".weight_scale", (0, 5120), (a, b), dst + ) + assert torch.equal(dst, t["down_s"][:, 40 * r : 40 * (r + 1)]) + finally: + reader.close() + + +def test_rank_column_slices_tile_the_logical_expert(quartets) -> None: + """The four down slices, side by side, are the stored tensor: no gap, overlap or reorder.""" + path, expected, t, _ = quartets + catalog = loader.SafetensorsCatalog(path) + spec = layout.QuantSpec( + "model.layers.3.mlp.experts.7.down_proj", (5120, 2560), 3, 7, "down_proj", 1 + ) + reader = loader.TileReader(1 << 16, pin=False) + parts, scales = [], [] + try: + for r in range(4): + dst = torch.empty(spec.local_packed_shape(r), dtype=torch.uint8) + loader.read_packed_into( + catalog, reader, spec.name + ".weight", (0, 5120), spec.packed_interval(r)[1:], dst + ) + parts.append(dst) + dst = torch.empty(spec.local_scale_shape(r), dtype=torch.uint8) + loader.read_packed_into( + catalog, + reader, + spec.name + ".weight_scale", + (0, 5120), + spec.scale_interval(r)[1:], + dst, + ) + scales.append(dst) + finally: + reader.close() + assert torch.equal(torch.cat(parts, 1), t["down_w"]) + assert torch.equal(torch.cat(scales, 1), t["down_s"]) + # Each rank's 640 logical columns decode from exactly its 320 bytes and 40 scales. + full = decode(t["down_w"][:8], t["down_s"][:8]) + for r in range(4): + assert torch.equal(decode(parts[r][:8], scales[r][:8]), full[:, 640 * r : 640 * (r + 1)]) + + +def test_scalars_read_exactly(quartets) -> None: + path, _, _, scalars = quartets + catalog = loader.SafetensorsCatalog(path) + reader = loader.TileReader(1 << 16, pin=False) + try: + got = loader.read_scalars(catalog, reader, list(scalars)) + finally: + reader.close() + assert {n: struct.pack(" None: + ok = torch.tensor([[0x00, 0x80, 0x01, 0x7E, 0x38]], dtype=torch.uint8) + loader.check_block_scales("ok", ok) + for bad in (0x7F, 0xFF, 0x81, 0xB8): + with pytest.raises(loader.CheckpointFormatError): + loader.check_block_scales("bad", torch.tensor([[0x38, bad]], dtype=torch.uint8)) + + +@pytest.mark.parametrize("how", ["missing", "extra", "dtype", "shape", "scalar_shape"]) +def test_validate_rejects_malformed_metadata(quartets, tmp_path, how) -> None: + path, expected, _, _ = quartets + catalog = loader.SafetensorsCatalog(path) + want = dict(expected) + name = "model.layers.3.mlp.experts.7.gate_proj.weight" + if how == "missing": + want["model.layers.3.mlp.experts.7.up_proj.weight"] = ("U8", (2560, 2560)) + elif how == "extra": + del want[name] + elif how == "dtype": + want[name] = ("I8", (2560, 2560)) + elif how == "shape": + want[name] = ("U8", (2560, 1280)) + else: + want["model.layers.3.mlp.experts.7.gate_proj.input_scale"] = ("F32", (1,)) + with pytest.raises(loader.CheckpointFormatError): + loader.validate(catalog, want) + + +def test_wrong_destination_raises(quartets) -> None: + path, _, _, _ = quartets + catalog = loader.SafetensorsCatalog(path) + reader = loader.TileReader(1 << 16, pin=False) + try: + with pytest.raises(ValueError): + loader.read_packed_into( + catalog, + reader, + "model.layers.3.mlp.experts.7.gate_proj.weight", + (0, 640), + None, + torch.empty(640, 2560, dtype=torch.int8), + ) + with pytest.raises(ValueError): + loader.read_packed_into( + catalog, + reader, + "model.layers.3.mlp.experts.7.gate_proj.weight", + (0, 640), + (0, 2561), + torch.empty(640, 2561, dtype=torch.uint8), + ) + finally: + reader.close() + + +# ------------------------------------------------------------------ CUDA: relayout and expert + + +def _sm103() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 3) + + +needs_sm103 = pytest.mark.skipif(not _sm103(), reason="the NVFP4 kernels are exercised on sm_103") + + +@needs_sm103 +@pytest.mark.parametrize("rows, cols", [(640, 320), (5120, 40), (256, 64)]) +def test_block_scale_relayout_matches_independent_offsets(rows, cols) -> None: + sf = _bytes(rows, cols, rows + cols).cuda() + out = torch.ops.trtllm.block_scale_interleave(sf).reshape(-1) + assert out.numel() == layout.swizzled_numel(rows, cols) + assert torch.equal(unswizzle(out, rows, cols), sf.cpu()) + + +def _e4m3_byte(x: torch.Tensor) -> torch.Tensor: + """FP32 -> E4M3 bytes, round to nearest even, saturating at 448 (torch's own cast).""" + return x.clamp(max=448.0).to(torch.float8_e4m3fn).view(torch.uint8) + + +def _e2m1_code(x: torch.Tensor) -> torch.Tensor: + """FP32 -> E2M1 codes, round to nearest even (ties to the even magnitude index), saturating.""" + mags = torch.tensor(_E2M1, dtype=torch.float32) + a = x.abs().clamp(max=6.0) + d = (a.unsqueeze(-1) - mags).abs() + best = d.min(-1, keepdim=True).values + ties = d == best + idx = torch.where(ties, torch.arange(8), 100) + even = torch.where(ties & (torch.arange(8) % 2 == 0), torch.arange(8), 100) + pick = torch.where(ties.sum(-1) > 1, even.min(-1).values, idx.min(-1).values) + return pick | ((x < 0) & (pick != 0)).long() * 8 + + +def _quantize(x: torch.Tensor, s: float) -> tuple: + """Documented static activation quantization, in the kernel's operation order: g = 1/s, + scale = e4m3(g * (amax * (1/6))), code = e2m1(x * (1 / (scale * (1/g)))). Returns codes + ``[M, K]``, block bytes ``[M, K/16]`` and the pre-rounding values.""" + one = torch.tensor(1.0, dtype=torch.float32) + g = one / torch.tensor(s, dtype=torch.float32) + blocks = x.float().cpu().view(x.shape[0], -1, 16) + sf = _e4m3_byte(g * (blocks.abs().amax(-1) * (one / 6.0))) + scale = _E4M3_TABLE[sf.long()].float() + out_scale = torch.where(scale > 0, one / (scale * (one / g)), torch.zeros_like(scale)) + scaled = (blocks * out_scale.unsqueeze(-1)).view(x.shape[0], -1) + return _e2m1_code(scaled), sf, scaled + + +#: Relative distance to an E2M1 midpoint inside which the kernel's approximate reciprocals may +#: round either way (the quantizer catalog entry's tie window). +_TIE = 2.0**-20 + + +def _near_tie(scaled: torch.Tensor) -> torch.Tensor: + mids = torch.tensor((0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)) + a = scaled.abs().unsqueeze(-1) + return ((a - mids).abs() <= _TIE * a).any(-1) + + +def _gemm(a_vals: torch.Tensor, w_vals: torch.Tensor, alpha: float) -> tuple: + acc = (a_vals @ w_vals.T).float() + out = (torch.tensor(alpha, dtype=torch.float32) * acc).to(torch.bfloat16) + abs_dot = (a_vals.abs() @ w_vals.abs().T).float() * abs(alpha) + return out, abs_dot + + +def _assert_gemm(out: torch.Tensor, a_codes, a_sf, w_packed, w_sf, alpha) -> None: + """Kernel GEMM output vs the exact product of the same quantized operands.""" + a_vals = _E2M1_TABLE[a_codes.cpu()].view(a_codes.shape[0], -1, 16) * _E4M3_TABLE[ + a_sf.cpu().long() + ].unsqueeze(-1) + ref, abs_dot = _gemm(a_vals.view(a_codes.shape[0], -1), decode(w_packed, w_sf), alpha) + r32 = ref.float() + ulp = torch.exp2(torch.frexp(r32.abs().clamp(min=2.0**-126))[1].float() - 8) + err = (out.cpu().float() - r32).abs() + assert bool((err <= ulp + 2.0**-20 * abs_dot).all()), float((err / ulp).max()) + + +@needs_sm103 +@pytest.mark.parametrize("shared_fc1", [True, False]) +def test_expert_provider_matches_independent_reference(shared_fc1) -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_nvfp4__sm_103__tp4.expert import ( # noqa: E501 + nvfp4_swiglu_expert, + ) + + hidden, inter, rows, n_experts, expert = 1024, 256, 7, 3, 2 + gen = torch.Generator().manual_seed(7) + + def proj(n, k, w2, s, seed): + # every code point including signed zero; block scales spread over normal E4M3 values, + # one all-zero block scale per projection + packed = _bytes(n, k // 2, seed) + sf = torch.randint(0x28, 0x48, (n, k // 16), dtype=torch.uint8, generator=gen) + sf[0, 0] = 0 + # every other expert slot holds different codes and block scales + bank = { + "weight": torch.stack( + [_bytes(n, k // 2, seed + 100 * (i + 1)) for i in range(n_experts)] + ), + "weight_scale": None, + "linear_scale": torch.stack([(sf + 1 + i).to(torch.uint8) for i in range(n_experts)]), + "input_scale": torch.full((n_experts,), s, dtype=torch.float32), + "weight_scale_2": torch.full((n_experts,), w2, dtype=torch.float32), + } + bank["weight"][expert] = packed + bank["linear_scale"][expert] = sf + bank["g"] = torch.ones(n_experts, dtype=torch.float32) / bank["input_scale"] + bank["alpha"] = bank["input_scale"] * bank["weight_scale_2"] + bank["weight_scale"] = torch.stack( + [ + torch.ops.trtllm.block_scale_interleave(bank["linear_scale"][i].cuda()) + .reshape(-1) + .cpu() + for i in range(n_experts) + ] + ) + return {key: v.cuda() for key, v in bank.items()} + + s_up = 0.0123 if shared_fc1 else 0.0171 + bank = { + "gate": proj(inter, hidden, 2.0**-9, 0.0123, 11), + "up": proj(inter, hidden, 3.0e-3, s_up, 12), + "down": proj(hidden, inter, 1.25e-3, 0.0625, 13), + } + x = (torch.randn(rows, hidden, generator=gen) * 0.7).to(torch.bfloat16) + x[0, :16] = 0 # one all-zero activation block + record = [] + y = nvfp4_swiglu_expert(x.cuda(), bank, expert, shared_fc1, record) + call = record[0] + assert call["expert"] == expert and y.dtype == torch.bfloat16 and y.shape == (rows, hidden) + assert (call["fc1_up"] is None) == shared_fc1 + + def check_quant(codes_packed, sf_flat, src, s): + k = src.shape[1] + codes = unpack(codes_packed.cpu()) + sf = unswizzle(sf_flat, src.shape[0], k // 16) + ref_codes, ref_sf, scaled = _quantize(src, s) + assert torch.equal(sf, ref_sf) + vals = _E2M1_TABLE[codes] * _E4M3_TABLE[sf.long()].repeat_interleave(16, 1) + ref_vals = _E2M1_TABLE[ref_codes] * _E4M3_TABLE[ref_sf.long()].repeat_interleave(16, 1) + differ = vals != ref_vals + assert not bool((differ & ~_near_tie(scaled)).any()) + assert int(differ.sum()) <= 2 + return codes, sf + + fc1_codes, fc1_sf = check_quant(call["fc1"], call["sf1"], x, 0.0123) + up_codes, up_sf = ( + (fc1_codes, fc1_sf) if shared_fc1 else check_quant(call["fc1_up"], call["sf1_up"], x, s_up) + ) + lin = {k: bank[k]["linear_scale"][expert] for k in bank} + _assert_gemm( + call["gate"], + fc1_codes, + fc1_sf, + bank["gate"]["weight"][expert], + lin["gate"], + 0.0123 * 2.0**-9, + ) + _assert_gemm( + call["up"], + up_codes, + up_sf, + bank["up"]["weight"][expert], + lin["up"], + float(torch.tensor(s_up, dtype=torch.float32) * torch.tensor(3.0e-3, dtype=torch.float32)), + ) + act = (torch.nn.functional.silu(call["gate"].float()) * call["up"].float()).to(torch.bfloat16) + assert torch.equal(act, call["act"]) + down_codes, down_sf = check_quant(call["fc2"], call["sf2"], call["act"].cpu(), 0.0625) + _assert_gemm( + y, + down_codes, + down_sf, + bank["down"]["weight"][expert], + lin["down"], + float( + torch.tensor(0.0625, dtype=torch.float32) * torch.tensor(1.25e-3, dtype=torch.float32) + ), + ) + # The provider used expert `expert`'s operands, not a neighbour's. + other = nvfp4_swiglu_expert(x.cuda(), bank, (expert + 1) % n_experts, shared_fc1) + assert not torch.equal(other, y) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/_dense_block_reference.py b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/_dense_block_reference.py new file mode 100644 index 000000000000..5578100d6cdb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/_dense_block_reference.py @@ -0,0 +1,501 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plain-Torch reference for one dense sliding-attention decoder block (block 0). + +An independent restatement of the source block semantics for generated-input +component replay: sandwich RMSNorms (fp32 statistics and affine, one cast), +separate bias-free Q/K/V/gate projections in [out, in] layout, per-head Q/K +RMSNorm rounded to the activation dtype before fp32 NeoX RoPE, GQA attention +with a causal window mask, per-head sigmoid output gate in the activation +dtype, bias-free output projection, fp32 SiLU(gate) * up cast once before the +down projection, and depth-scaled residual adds in the activation dtype. + +Only torch is imported. `Variant` selects documented equivalent formulations +(valid alternatives) or deliberately wrong ones (fault witnesses); the default +is the source formulation. +""" + +from __future__ import annotations + +import contextlib +import math +from collections import OrderedDict +from dataclasses import dataclass +from typing import Callable, Iterator + +import torch +import torch.nn.functional as F + + +@dataclass(frozen=True) +class Dims: + hidden: int = 5120 + intermediate: int = 20480 + heads: int = 48 + kv_heads: int = 8 + head_dim: int = 128 + eps: float = 1e-6 + rope_theta: float = 10000.0 + rope_factor: float = 1.0 + window: int = 512 + layer_id: int = 0 + max_len: int = 129 + + +BLOCK0 = Dims() +SEEDS = (17, 41) +LENGTHS = (1, 17, 129) + +STATE_RECIPE = { + "construction": "torch.manual_seed(1300 + seed + layer_id); fp32 kaiming_uniform_(a=sqrt(5)) on CPU for the " + "eight projection matrices in module construction order (q, k, v, gate, o, mlp gate, up, down)", + "overrides": "torch.Generator().manual_seed(10000 + seed) in state_dict order: *norm.weight = 1 + 0.25 * N(0, 1); " + "self_attn.gate_bias = 0.75 * N(0, 1)", + "dtype": "every tensor cast once to bfloat16", + "input": "torch.Generator().manual_seed(1200 + seed): 2 * N(0, 1) of shape [1, max_len, hidden], cast once to " + "bfloat16; length-T cases use the first T rows; positions arange(T)", +} + + +def depth_coefs(layer_id: int, is_attention: bool) -> tuple[float, float]: + one_based = layer_id + 1 + num_adds = 2 * one_based if is_attention else 2 * one_based + 1 + denom = math.sqrt(num_adds) + return 1.0 / denom, math.sqrt(num_adds - 1) / denom + + +def param_shapes(d: Dims = BLOCK0) -> "OrderedDict[str, tuple[int, ...]]": + """Block parameters in source state_dict order.""" + q, kv = d.heads * d.head_dim, d.kv_heads * d.head_dim + return OrderedDict( + [ + ("input_layernorm.weight", (d.hidden,)), + ("self_attn.gate_bias", (d.heads,)), + ("self_attn.q_proj.weight", (q, d.hidden)), + ("self_attn.k_proj.weight", (kv, d.hidden)), + ("self_attn.v_proj.weight", (kv, d.hidden)), + ("self_attn.gate_proj.weight", (q, d.hidden)), + ("self_attn.o_proj.weight", (d.hidden, q)), + ("self_attn.q_norm.weight", (d.head_dim,)), + ("self_attn.k_norm.weight", (d.head_dim,)), + ("pre_feedforward_layernorm.weight", (d.hidden,)), + ("mlp.gate_proj.weight", (d.intermediate, d.hidden)), + ("mlp.up_proj.weight", (d.intermediate, d.hidden)), + ("mlp.down_proj.weight", (d.hidden, d.intermediate)), + ("post_attention_layernorm.weight", (d.hidden,)), + ("post_feedforward_layernorm.weight", (d.hidden,)), + ] + ) + + +INIT_ORDER = ( + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", + "self_attn.gate_proj.weight", + "self_attn.o_proj.weight", + "mlp.gate_proj.weight", + "mlp.up_proj.weight", + "mlp.down_proj.weight", +) + + +def generate_state(seed: int, d: Dims = BLOCK0) -> "OrderedDict[str, torch.Tensor]": + """Deterministic bf16 block state on CPU (see STATE_RECIPE).""" + shapes = param_shapes(d) + torch.manual_seed(1300 + seed + d.layer_id) + matrices = {} + for name in INIT_ORDER: + weight = torch.empty(shapes[name], dtype=torch.float32) + torch.nn.init.kaiming_uniform_(weight, a=math.sqrt(5)) + matrices[name] = weight + gen = torch.Generator().manual_seed(10000 + seed) + state: "OrderedDict[str, torch.Tensor]" = OrderedDict() + for name, shape in shapes.items(): + if name.endswith("norm.weight"): + value = 1.0 + 0.25 * torch.randn(shape, generator=gen) + elif name == "self_attn.gate_bias": + value = 0.75 * torch.randn(shape, generator=gen) + else: + value = matrices[name] + state[name] = value.to(torch.bfloat16).contiguous() + return state + + +def generate_input(seed: int, d: Dims = BLOCK0) -> torch.Tensor: + gen = torch.Generator().manual_seed(1200 + seed) + return (2.0 * torch.randn(1, d.max_len, d.hidden, generator=gen)).to(torch.bfloat16) + + +def state_problems(state: "dict[str, torch.Tensor]", d: Dims = BLOCK0) -> list[str]: + """Exact accounting plus the nontrivial-state properties the comparisons rely on.""" + problems = [] + shapes = param_shapes(d) + if list(state) != list(shapes): + problems.append(f"names/order {list(state)} != {list(shapes)}") + for name, shape in shapes.items(): + t = state.get(name) + if t is None: + continue + if tuple(t.shape) != shape or t.dtype != torch.bfloat16: + problems.append(f"{name}: {tuple(t.shape)} {t.dtype}, expected {shape} bfloat16") + if not bool(torch.isfinite(t.float()).all()): + problems.append(f"{name}: nonfinite") + norms = [n for n in shapes if n.endswith("norm.weight")] + for name in norms: + if name in state and bool((state[name].float() == 1.0).all()): + problems.append(f"{name}: unit affine") + for i, a in enumerate(norms): + for b in norms[i + 1 :]: + if ( + a in state + and b in state + and state[a].shape == state[b].shape + and torch.equal(state[a], state[b]) + ): + problems.append(f"{a} and {b} are identical") + bias = state.get("self_attn.gate_bias") + if bias is not None: + if not bool((bias.float() != 0).any()): + problems.append("gate_bias is all zero") + if bool((bias == bias[0]).all()): + problems.append("gate_bias does not vary across heads") + return problems + + +@dataclass(frozen=True) +class Variant: + """Source formulation by default; see module docstring.""" + + # Valid alternatives (same mathematics, different evaluation). + sdpa_backend: str | None = ( + None # None: default dispatch; MATH | EFFICIENT_ATTENTION | FLASH_ATTENTION | CUDNN_ATTENTION + ) + sdpa_form: str = "mask" # mask: explicit window mask; is_causal: causal flag (equal for T <= window); enable_gqa + gemm: str = "batched" # batched | per_row | fused | matmul_t + rmsnorm: str = "source" # source | alt_order | chunked + gate: str = "source" # source | fp32_sigmoid_then_cast + residual: str = "source" # source | torch_add + silu: str = "source" # source: F.silu | x_sigmoid: x * sigmoid(x) | exp_form: x / (1 + exp(-x)); all fp32 + # Fault witnesses (wrong semantics). + attn_coefs: tuple[float, float] | None = None + mlp_coefs: tuple[float, float] | None = None + residual_fp32_fused: bool = False + omit_post_attention_norm: bool = False + swap_post_attention_and_pre_ffn_norms: bool = False + rmsnorm_round_before_affine: bool = False + softmax_scale: float | None = None + gqa_tiled: bool = False + noncausal: bool = False + omit_k_norm: bool = False + swap_qk_norm_weights: bool = False + gate_bias: str = "source" # source | omitted | uniform_mean | heads_rolled + swiglu_bf16_activation: bool = False + rope_interleaved: bool = False + rope_positions_offset: int = 0 + normalize_v: bool = False + qk_norm_fp32_into_rope: bool = False + + +SOURCE = Variant() + +_SDPA = { + "MATH": "MATH", + "EFFICIENT_ATTENTION": "EFFICIENT_ATTENTION", + "FLASH_ATTENTION": "FLASH_ATTENTION", + "CUDNN_ATTENTION": "CUDNN_ATTENTION", +} + + +@contextlib.contextmanager +def sdpa_backend(name: str | None) -> Iterator[None]: + if name is None: + yield + return + from torch.nn.attention import SDPBackend, sdpa_kernel + + with sdpa_kernel([getattr(SDPBackend, _SDPA[name])]): + yield + + +Record = Callable[[str, torch.Tensor], None] + + +def _noop(name: str, value: torch.Tensor) -> None: + del name, value + + +def linear(x: torch.Tensor, weight: torch.Tensor, v: Variant) -> torch.Tensor: + if v.gemm == "per_row": + rows = [F.linear(x[..., i : i + 1, :], weight) for i in range(x.shape[-2])] + return torch.cat(rows, dim=-2) + if v.gemm == "matmul_t": + return torch.matmul(x, weight.t()) + return F.linear(x, weight) + + +def rms_norm( + x: torch.Tensor, weight: torch.Tensor, eps: float, v: Variant = SOURCE +) -> torch.Tensor: + in_dtype = x.dtype + x32 = x.float() + if v.rmsnorm == "alt_order": + normed = x32 / torch.sqrt(x32.pow(2).sum(-1, keepdim=True) / x32.shape[-1] + eps) + elif v.rmsnorm == "chunked": + n = x32.shape[-1] + sq = x32.pow(2).unflatten(-1, (n // 32, 32)).sum(-1).sum(-1, keepdim=True) + normed = x32 * (sq / n + eps).rsqrt() + else: + rms = x32.pow(2).mean(dim=-1, keepdim=True).add(eps).rsqrt() + normed = x32 * rms + if v.rmsnorm_round_before_affine: + return normed.to(in_dtype) * weight + return (normed * weight).to(in_dtype) + + +def rope_cos_sin(positions: torch.Tensor, d: Dims = BLOCK0) -> tuple[torch.Tensor, torch.Tensor]: + exponent = torch.arange(0, d.head_dim, 2, dtype=torch.float32, device=positions.device) + inv_freq = 1.0 / (d.rope_theta ** (exponent / d.head_dim)) + inv_freq = inv_freq / d.rope_factor + freqs = positions.float().unsqueeze(-1) * inv_freq + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos(), emb.sin() + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def apply_rope( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, v: Variant = SOURCE +) -> tuple[torch.Tensor, torch.Tensor]: + """q, k: (B, T, heads, head_dim); cos, sin: (B, T, head_dim) fp32.""" + q_dtype, k_dtype = q.dtype, k.dtype + cos = cos.unsqueeze(2).float() + sin = sin.unsqueeze(2).float() + q, k = q.float(), k.float() + if v.rope_interleaved: + + def rot(x: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + c, s = cos[..., :half], sin[..., :half] + even, odd = x[..., 0::2], x[..., 1::2] + return torch.stack((even * c - odd * s, even * s + odd * c), dim=-1).flatten(-2) + + return rot(q).to(q_dtype), rot(k).to(k_dtype) + q_out = q * cos + _rotate_half(q) * sin + k_out = k * cos + _rotate_half(k) * sin + return q_out.to(q_dtype), k_out.to(k_dtype) + + +def causal_window_mask(length: int, window: int | None, device: torch.device) -> torch.Tensor: + q = torch.arange(length, device=device).unsqueeze(1) + k = torch.arange(length, device=device).unsqueeze(0) + mask = q >= k + if window is not None: + mask = mask & ((q - k) < window) + return mask.unsqueeze(0).unsqueeze(0) + + +def attention_core( + q: torch.Tensor, k: torch.Tensor, val: torch.Tensor, d: Dims = BLOCK0, v: Variant = SOURCE +) -> torch.Tensor: + """q: (B, T, heads, D) post-RoPE; k, val: (B, T, kv_heads, D). Returns (B, T, heads, D).""" + length = q.shape[1] + q, k, val = q.transpose(1, 2), k.transpose(1, 2), val.transpose(1, 2) + repeat = d.heads // d.kv_heads + gqa_native = v.sdpa_form == "enable_gqa" + if not gqa_native and repeat != 1: + if v.gqa_tiled: + k, val = k.repeat(1, repeat, 1, 1), val.repeat(1, repeat, 1, 1) + else: + k, val = k.repeat_interleave(repeat, dim=1), val.repeat_interleave(repeat, dim=1) + scale = d.head_dim**-0.5 if v.softmax_scale is None else v.softmax_scale + kwargs = {"dropout_p": 0.0, "scale": scale} + if gqa_native: + kwargs["enable_gqa"] = True + with sdpa_backend(v.sdpa_backend): + if v.noncausal: + out = F.scaled_dot_product_attention(q, k, val, **kwargs) + elif v.sdpa_form == "is_causal": + out = F.scaled_dot_product_attention(q, k, val, is_causal=True, **kwargs) + else: + mask = causal_window_mask(length, d.window, q.device) + out = F.scaled_dot_product_attention(q, k, val, attn_mask=mask, **kwargs) + return out.transpose(1, 2) + + +def output_gate( + attn: torch.Tensor, gate: torch.Tensor, gate_bias: torch.Tensor, v: Variant = SOURCE +) -> torch.Tensor: + """attn, gate: (B, T, heads, D) activation dtype; gate_bias: (heads,).""" + if v.gate_bias == "omitted": + gate_bias = torch.zeros_like(gate_bias) + elif v.gate_bias == "uniform_mean": + gate_bias = torch.full_like(gate_bias, float(gate_bias.float().mean())) + elif v.gate_bias == "heads_rolled": + gate_bias = gate_bias.roll(1) + out_dtype = attn.dtype + biased = gate + gate_bias[None, None, :, None] + if v.gate == "fp32_sigmoid_then_cast": + return attn * torch.sigmoid(biased.float()).to(out_dtype) + return (attn * torch.sigmoid(biased)).to(out_dtype) + + +def attention( + x: torch.Tensor, + positions: torch.Tensor, + s: "dict[str, torch.Tensor]", + d: Dims = BLOCK0, + v: Variant = SOURCE, + record: Record = _noop, + fused: "dict[str, torch.Tensor] | None" = None, +) -> torch.Tensor: + b, t, _ = x.shape + qw, kvw = d.heads * d.head_dim, d.kv_heads * d.head_dim + if v.gemm == "fused": + assert fused is not None + q, k, val, gate = ( + y.contiguous() for y in F.linear(x, fused["qkvg"]).split([qw, kvw, kvw, qw], dim=-1) + ) + else: + q = linear(x, s["self_attn.q_proj.weight"], v) + k = linear(x, s["self_attn.k_proj.weight"], v) + val = linear(x, s["self_attn.v_proj.weight"], v) + gate = linear(x, s["self_attn.gate_proj.weight"], v) + for name, value in (("q_proj", q), ("k_proj", k), ("v_proj", val), ("gate_proj", gate)): + record(name, value) + q = q.view(b, t, d.heads, d.head_dim) + k = k.view(b, t, d.kv_heads, d.head_dim) + val = val.view(b, t, d.kv_heads, d.head_dim) + gate = gate.view(b, t, d.heads, d.head_dim) + qn_w, kn_w = s["self_attn.q_norm.weight"], s["self_attn.k_norm.weight"] + if v.swap_qk_norm_weights: + qn_w, kn_w = kn_w, qn_w + if v.qk_norm_fp32_into_rope: + q32 = q.float() + q = (q32 * q32.pow(2).mean(-1, keepdim=True).add(d.eps).rsqrt()) * qn_w + k32 = k.float() + k = (k32 * k32.pow(2).mean(-1, keepdim=True).add(d.eps).rsqrt()) * kn_w + qn, kn = q.to(torch.bfloat16), k.to(torch.bfloat16) + else: + q = rms_norm(q, qn_w, d.eps, v) + if not v.omit_k_norm: + k = rms_norm(k, kn_w, d.eps, v) + qn, kn = q, k + if v.normalize_v: + val = rms_norm(val, kn_w, d.eps, v) + record("q_norm", qn) + record("k_norm", kn) + cos, sin = rope_cos_sin(positions + v.rope_positions_offset, d) + q, k = apply_rope(q, k, cos, sin, v) + if v.qk_norm_fp32_into_rope: + q, k = q.to(torch.bfloat16), k.to(torch.bfloat16) + record("q_rope", q) + record("k_rope", k) + attn = attention_core(q, k, val, d, v) + record("attention_core", attn) + attn = output_gate(attn, gate, s["self_attn.gate_bias"], v) + attn = attn.reshape(b, t, qw) + record("gated", attn) + out = linear(attn, s["self_attn.o_proj.weight"], v) + record("attention_output", out) + return out + + +def silu_and_mul(gate: torch.Tensor, up: torch.Tensor, v: Variant = SOURCE) -> torch.Tensor: + """fp32 SiLU(gate) * up, cast once to the activation dtype.""" + if v.swiglu_bf16_activation: + return F.silu(gate) * up + g, u = gate.float(), up.float() + if v.silu == "x_sigmoid": + act = g * torch.sigmoid(g) + elif v.silu == "exp_form": + act = g / (1.0 + torch.exp(-g)) + else: + act = F.silu(g) + return (act * u).to(gate.dtype) + + +def swiglu( + x: torch.Tensor, + s: "dict[str, torch.Tensor]", + d: Dims = BLOCK0, + v: Variant = SOURCE, + record: Record = _noop, + fused: "dict[str, torch.Tensor] | None" = None, +) -> torch.Tensor: + if v.gemm == "fused": + assert fused is not None + gate, up = ( + y.contiguous() + for y in F.linear(x, fused["gate_up"]).split([d.intermediate] * 2, dim=-1) + ) + else: + gate = linear(x, s["mlp.gate_proj.weight"], v) + up = linear(x, s["mlp.up_proj.weight"], v) + record("mlp_gate", gate) + record("mlp_up", up) + act = silu_and_mul(gate, up, v) + record("mlp_act", act) + out = linear(act, s["mlp.down_proj.weight"], v) + record("mlp_output", out) + return out + + +def residual_add( + x: torch.Tensor, residual: torch.Tensor, scales: tuple[float, float], v: Variant = SOURCE +) -> torch.Tensor: + x_scale, r_scale = scales + if v.residual_fp32_fused: + return (x.float() * x_scale + residual.float() * r_scale).to(x.dtype) + if v.residual == "torch_add": + return torch.add(x * x_scale, residual * r_scale) + return x * x_scale + residual * r_scale + + +def fused_weights(s: "dict[str, torch.Tensor]") -> "dict[str, torch.Tensor]": + return { + "qkvg": torch.cat( + [s[f"self_attn.{n}_proj.weight"] for n in ("q", "k", "v", "gate")], dim=0 + ).contiguous(), + "gate_up": torch.cat( + [s["mlp.gate_proj.weight"], s["mlp.up_proj.weight"]], dim=0 + ).contiguous(), + } + + +def block( + x: torch.Tensor, + positions: torch.Tensor, + s: "dict[str, torch.Tensor]", + d: Dims = BLOCK0, + v: Variant = SOURCE, + record: Record = _noop, + fused: "dict[str, torch.Tensor] | None" = None, +) -> torch.Tensor: + """x: (1, T, hidden) bf16; positions: (1, T) int64. Returns (1, T, hidden) bf16.""" + attn_scales = v.attn_coefs or depth_coefs(d.layer_id, True) + mlp_scales = v.mlp_coefs or depth_coefs(d.layer_id, False) + post_attn_w = s["post_attention_layernorm.weight"] + pre_ffn_w = s["pre_feedforward_layernorm.weight"] + if v.swap_post_attention_and_pre_ffn_norms: + post_attn_w, pre_ffn_w = pre_ffn_w, post_attn_w + residual = x + h = rms_norm(x, s["input_layernorm.weight"], d.eps, v) + record("input_layernorm", h) + h = attention(h, positions, s, d, v, record, fused) + if not v.omit_post_attention_norm: + h = rms_norm(h, post_attn_w, d.eps, v) + record("post_attention_layernorm", h) + hidden = residual_add(h, residual, attn_scales, v) + record("attention_residual", hidden) + h = rms_norm(hidden, pre_ffn_w, d.eps, v) + record("pre_feedforward_layernorm", h) + h = swiglu(h, s, d, v, record, fused) + h = rms_norm(h, s["post_feedforward_layernorm.weight"], d.eps, v) + record("post_feedforward_layernorm", h) + out = residual_add(h, hidden, mlp_scales, v) + record("block_output", out) + return out diff --git a/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/test_dense_block.py b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/test_dense_block.py new file mode 100644 index 000000000000..618859de29d2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/test_dense_block.py @@ -0,0 +1,1123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the Beam dense block 0 component: bf16 / sm_103 / tp1. + +Configuration under test, fixed below and nowhere else: block 0 (sliding attention, +dense MLP) at hidden 5120, intermediate 20480 and 48 q / 8 kv heads of 128; bf16 +weights and activations with the source model's fp32 steps; one sm_103 GPU; eager, +batch 1, one fresh context per call; generated weights; seeds {17, 41} x lengths +{1, 17, 129}. + +What is compared. ``_dense_block_reference.py`` restates the source block in plain +torch and generates its state and input deterministically on the CPU. For every case +the component's final output is compared with the reference's, and every operation of +the block is replayed on the reference's own input to that operation, so a local error +cannot hide behind the sandwich norms. Each compared tensor is held to one rule in +`RULES`. + +Invocation (fp32 steps need the TF32 overrides removed before torch is imported):: + + 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 + +Optional environment: + +* ``BEAM_DENSE_BLOCK_SOURCE_CAPTURES``: a directory of ``seed/T.pt`` files, each a + dict of tensors captured from the source block under the reference's record names + (plus ``input`` and ``positions``). The final output and every replay are then also + compared against the captures, and the generated input must equal theirs bit for bit. +* ``BEAM_DENSE_BLOCK_EVIDENCE``: a directory that receives one JSON record per test case: + every comparison with its rule, the op calls and CUDA kernels of the forward, and the + device and software it ran on; and ``fp32-session.json``, the fp32 GEMM checks made + before and after the module's tests. + +Dependencies: a TensorRT-LLM build for sm_103 with flashinfer importable. No model files. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import json +import math +import os +import platform +import re +import subprocess +import sys +import types +from pathlib import Path + +import pytest +import torch + +import tensorrt_llm +import tensorrt_llm._torch.models.modeling_beam as builtin_beam +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501 + flashinfer_silu_and_mul, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_neox_rope import ( + beam_neox_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_output_gate import ( + beam_output_gate, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + thop_attention, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.residual.beam_depth_residual import ( + beam_depth_residual, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam import _attention +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1 import ( # noqa: E501 + block as block_module, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1.block import ( # noqa: E501 + BeamDenseBlockBf16Sm103Tp1, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping + +__extra_import_path__ = [".", "../../.."] +import _dense_block_reference as reference # noqa: E402 +from _validating import validating # noqa: E402 + +# ----------------------------------------------------------------------------- the fixed configuration + +#: The block as the checkpoint configures it. Only these fields are read. +BLOCK_CONFIG = types.MappingProxyType( + dict( + hidden_size=5120, + intermediate_size=20480, + num_attention_heads=48, + num_key_value_heads=8, + head_dim=128, + rms_norm_eps=1e-6, + hidden_act="silu", + attention_bias=False, + attention_dropout=0.0, + sliding_window=512, + layer_types=("sliding_attention",), + mlp_layer_types=("dense",), + rope_parameters={ + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + }, + ) +) +LAYER_IDX = 0 +DTYPE = torch.bfloat16 +SM = (10, 3) +SEEDS = (17, 41) +LENGTHS = (1, 17, 129) +CASES = tuple((seed, length) for seed in SEEDS for length in LENGTHS) +CASE_IDS = [f"seed{seed}-T{length}" for seed, length in CASES] +#: Eager, batch 1, one fresh context: no CUDA graph, no overlap scheduler, no cached prefix. +EXECUTION = types.MappingProxyType( + dict(eager=True, batch_size=1, fresh_context=True, cuda_graph=False, overlap_scheduler=False) +) +#: The KV cache the attention op appends to: one bf16 layer, 32-token pages, 512 tokens per +#: sequence (the window), one request. +KV_CACHE = types.MappingProxyType(dict(tokens_per_block=32, max_seq_len=512, num_blocks=16)) +REQUEST_ID = 0 +ROPE_THETA = 10000.0 + +# ----------------------------------------------------------------------------- comparison rules + +_RTOL = 2.0**-7 + +#: (atol, rtol, max fraction of unequal elements or None) per kind of compared tensor. An +#: element passes when |a - r| <= atol + rtol * |r| with both widened to fp64; every element +#: of every tensor of every case must pass, and no NaN or Inf is accepted anywhere. +#: +#: Derivation: the largest atol each kind needed at rtol 2^-7 across valid evaluations of +#: the source block on these exact cases (other GEMM kernels, reduction orders, attention +#: backends and Torch builds) was 0 for the norms, RoPE, gate, residual adds and SiLU +#: product (with at most 5.7e-5 of elements one ulp apart, hence the 0.1% unequal bound), +#: 5.0e-6 for a single projection, 1.6e-4 for the dense SwiGLU output, 7.8e-4 for the +#: attention module and 0.0375 for the whole block. The bands below sit 100x, 4.9x, 3.8x +#: and 4.0x above those last four, and a zeroed final residual coefficient needs an atol +#: of 4.8. +RULES = types.MappingProxyType( + { + "rmsnorm": (1e-6, _RTOL, 1e-3), + "rope": (1e-6, _RTOL, 1e-3), + "gate": (1e-6, _RTOL, 1e-3), + "residual": (1e-6, _RTOL, 1e-3), + "swiglu_activation": (1e-6, _RTOL, 1e-3), + "projection": (5e-4, _RTOL, None), + "swiglu_output": (8e-4, _RTOL, None), + "attention_output": (3e-3, _RTOL, None), + "block_output": (0.15, _RTOL, None), + } +) + +#: The forward's op calls, in order: the composition this component claims to run. +EXPECTED_CALLS = ( + "FlashinferRmsnorm", + "CublasMm", + "CublasMm", + "CublasMm", + "CublasMm", + "FlashinferRmsnorm", + "FlashinferRmsnorm", + "BeamNeoxRope", + "ThopAttention", + "BeamOutputGate", + "CublasMm", + "FlashinferRmsnorm", + "BeamDepthResidual", + "FlashinferRmsnorm", + "CublasMm", + "CublasMm", + "FlashinferSiluAndMul", + "CublasMm", + "FlashinferRmsnorm", + "BeamDepthResidual", +) +OPS = ( + flashinfer_rmsnorm, + cublas_mm, + beam_neox_rope, + thop_attention, + beam_output_gate, + beam_depth_residual, + flashinfer_silu_and_mul, +) + +_SDPA_OPS = re.compile(r"scaled_dot_product|_flash_attention|_efficient_attention|_cudnn_attention") +_ATTENTION_KERNEL = re.compile(r"(?i)fmha|attention") + +assert torch.cuda.is_available(), "the Beam dense block requires a CUDA device" + + +def block_config(**overrides) -> types.SimpleNamespace: + fields = dict(BLOCK_CONFIG) + fields["layer_types"] = list(fields["layer_types"]) + fields["mlp_layer_types"] = list(fields["mlp_layer_types"]) + fields["rope_parameters"] = { + k: (dict(v) if v else v) for k, v in fields["rope_parameters"].items() + } + fields.update(overrides) + return types.SimpleNamespace(**fields) + + +# ----------------------------------------------------------------------------- precision and environment + + +def _fp32_gemm_probe() -> float: + """Worst relative error of an fp32 GEMM against fp64; TF32 substitution shows as ~1e-3.""" + gen = torch.Generator().manual_seed(1234) + x = torch.randn(64, 5120, generator=gen).to(torch.bfloat16).float().cuda() + w = (torch.randn(512, 5120, generator=gen) * 5120**-0.5).float().cuda() + y = torch.nn.functional.linear(x, w) + ref = x.double() @ w.double().t() + return float((y.double() - ref).abs().max()) / float(ref.pow(2).mean().sqrt()) + + +def _math_state() -> dict: + return { + "float32_matmul_precision": torch.get_float32_matmul_precision(), + "cuda_matmul_allow_tf32": bool(torch.backends.cuda.matmul.allow_tf32), + "cudnn_allow_tf32": bool(torch.backends.cudnn.allow_tf32), + } + + +@pytest.fixture(scope="module", autouse=True) +def ieee_fp32(): + """The source's fp32 steps are IEEE fp32: refuse a process that substitutes TF32, both + before the module's tests and after them.""" + for var in ("TORCH_ALLOW_TF32_CUBLAS_OVERRIDE", "NVIDIA_TF32_OVERRIDE"): + assert var not in os.environ, ( + f"{var} is set and substitutes TF32 for fp32 GEMMs; run with `env -u {var}`" + ) + saved = ( + torch.get_float32_matmul_precision(), + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + ) + torch.set_float32_matmul_precision("highest") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + before = _fp32_gemm_probe() + assert before <= 1e-5, f"fp32 GEMM is not IEEE fp32 (relative error {before:.3g})" + state = {"math": _math_state(), "fp32_gemm_probe_before": before} + yield state + after = _fp32_gemm_probe() + state.update(math_after=_math_state(), fp32_gemm_probe_after=after) + _write_evidence("fp32-session", {"fixture": "ieee_fp32", "math": state}) + precision, allow_matmul, allow_cudnn = saved + torch.set_float32_matmul_precision(precision) + torch.backends.cuda.matmul.allow_tf32 = allow_matmul + torch.backends.cudnn.allow_tf32 = allow_cudnn + assert after <= 1e-5, f"fp32 GEMM stopped being IEEE fp32 (relative error {after:.3g})" + + +def _git_revision() -> str | None: + root = Path(tensorrt_llm.__file__).resolve().parents[1] + try: + out = subprocess.run( + ["git", "-C", str(root), "rev-parse", "HEAD"], + capture_output=True, + text=True, + timeout=30, + ) + except (OSError, subprocess.SubprocessError): + return None + return out.stdout.strip() or None + + +def _environment(ieee: dict) -> dict: + device = torch.cuda.current_device() + return { + "device_name": torch.cuda.get_device_name(device), + "device_capability": list(torch.cuda.get_device_capability(device)), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "tensorrt_llm": tensorrt_llm.__version__, + "tensorrt_llm_file": tensorrt_llm.__file__, + "source_revision": _git_revision(), + "python": sys.version.split()[0], + "machine": platform.machine(), + "argv": sys.argv, + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "slurm_step_id": os.environ.get("SLURM_STEP_ID"), + "execution": dict(EXECUTION), + "math": ieee, + } + + +def _write_evidence(name: str, record: dict) -> None: + directory = os.environ.get("BEAM_DENSE_BLOCK_EVIDENCE") + if not directory: + return + path = Path(directory) + path.mkdir(parents=True, exist_ok=True) + (path / f"{name}.json").write_text(json.dumps(record, indent=1, default=str) + "\n") + + +def tensor_sha256(t: torch.Tensor) -> str: + flat = t.detach().contiguous().cpu().reshape(-1) + h = hashlib.sha256() + h.update(f"{t.dtype}|{tuple(t.shape)}|".encode()) + h.update(flat.view(torch.uint8).numpy().tobytes()) + return h.hexdigest() + + +# ----------------------------------------------------------------------------- comparison + + +def compare(name: str, actual: torch.Tensor, expected: torch.Tensor, rule: str | None) -> dict: + """Score one tensor against one rule; `rule=None` measures without judging.""" + result = { + "tensor": name, + "rule": rule, + "shape": list(expected.shape), + "dtype": str(actual.dtype), + } + problems = [] + if tuple(actual.shape) != tuple(expected.shape): + problems.append(f"shape {tuple(actual.shape)} != {tuple(expected.shape)}") + if actual.dtype != DTYPE or expected.dtype != DTYPE: + problems.append(f"dtypes {actual.dtype}/{expected.dtype}, need {DTYPE}") + if problems: + result.update(passed=False, problems=problems) + return result + a, r = actual.double(), expected.double() + if not bool(torch.isfinite(a).all()) or not bool(torch.isfinite(r).all()): + result.update(passed=False, problems=["non-finite values"]) + return result + diff = (a - r).abs() + numel = diff.numel() + worst = int(diff.argmax()) + result.update( + numel=numel, + max_abs=float(diff.max()), + mean_abs=float(diff.mean()), + worst_index=[int(i) for i in torch.unravel_index(torch.tensor(worst), diff.shape)], + worst_actual=float(a.reshape(-1)[worst]), + worst_expected=float(r.reshape(-1)[worst]), + unequal=int((a != r).sum()), + ) + if rule is None: + result["passed"] = None + return result + atol, rtol, max_unequal = RULES[rule] + violations = int((diff > atol + rtol * r.abs()).sum()) + result.update(atol=atol, rtol=rtol, violations=violations) + if violations: + problems.append( + f"{violations} of {numel} elements outside atol {atol} + rtol {rtol} * |ref|" + ) + if max_unequal is not None: + allowed = math.floor(max_unequal * numel) + result["max_unequal"] = allowed + if result["unequal"] > allowed: + problems.append(f"{result['unequal']} unequal elements, at most {allowed} allowed") + result.update(passed=not problems, problems=problems) + return result + + +def assert_all_pass(results: list[dict]) -> None: + failed = [r for r in results if r["passed"] is False] + assert not failed, "\n".join( + f"{r['tensor']} [{r['rule']}]: {'; '.join(r['problems'])}" for r in failed + ) + + +# ----------------------------------------------------------------------------- state, inputs, fresh context + +_STATE_CACHE: dict[int, dict] = {} + + +def generated(seed: int) -> dict: + """The seed's bf16 state and [1, 129, hidden] input, generated once on the CPU.""" + if seed not in _STATE_CACHE: + _STATE_CACHE.clear() # one seed's 0.8 GB of weights at a time + state = reference.generate_state(seed) + assert not reference.state_problems(state), reference.state_problems(state) + _STATE_CACHE[seed] = { + "state_cpu": state, + "state": {name: t.cuda() for name, t in state.items()}, + "input": reference.generate_input(seed).cuda(), + "block": None, + } + return _STATE_CACHE[seed] + + +def loaded_block(seed: int) -> BeamDenseBlockBf16Sm103Tp1: + entry = generated(seed) + if entry["block"] is None: + blk = BeamDenseBlockBf16Sm103Tp1(block_config(), LAYER_IDX, device="cuda") + blk.load_weights(entry["state"]) + entry["block"] = blk + return entry["block"] + + +@contextlib.contextmanager +def fresh_context(num_tokens: int, cached_tokens: int = 0, num_requests: int = 1, **mgr_overrides): + """A single-layer KVCacheManagerV2 and prepared TrtllmAttentionMetadata for one new request.""" + mgr_args = dict( + num_layers=1, + num_kv_heads=BLOCK_CONFIG["num_key_value_heads"], + head_dim=BLOCK_CONFIG["head_dim"], + tokens_per_block=KV_CACHE["tokens_per_block"], + max_seq_len=KV_CACHE["max_seq_len"], + max_batch_size=num_requests, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.BF16, + ) + mgr_args.update(mgr_overrides) + mgr = KVCacheManagerV2( + KvCacheConfig( + max_tokens=KV_CACHE["num_blocks"] * KV_CACHE["tokens_per_block"] * num_requests, + enable_block_reuse=False, + ), + CacheTypeCpp.SELF, + **mgr_args, + ) + try: + request_ids = [REQUEST_ID + i for i in range(num_requests)] + mgr.add_dummy_requests(request_ids, token_nums=[cached_tokens + num_tokens] * num_requests) + md = TrtllmAttentionMetadata( + max_num_requests=num_requests, + max_num_tokens=block_module.MAX_TOKENS * num_requests, + kv_cache_manager=mgr, + num_contexts=num_requests, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=[cached_tokens] * num_requests + ), + seq_lens=torch.tensor([num_tokens] * num_requests, dtype=torch.int), + request_ids=request_ids, + prompt_lens=[cached_tokens + num_tokens] * num_requests, + ) + md.prepare() + yield md + finally: + mgr.shutdown() + + +def case_inputs(seed: int, length: int) -> tuple[torch.Tensor, torch.Tensor]: + x = generated(seed)["input"][:, :length].contiguous() + positions = torch.arange(length, device="cuda").view(1, length) + return x, positions + + +def reference_record(seed: int, length: int) -> dict: + """The reference block's output and every named intermediate, on the GPU.""" + x, positions = case_inputs(seed, length) + record = {"input": x, "positions": positions} + + def keep(name: str, value: torch.Tensor) -> None: + record[name] = value + + with torch.no_grad(): + out = reference.block(x, positions, generated(seed)["state"], record=keep) + assert torch.equal(out, record["block_output"]) + return record + + +def source_captures(seed: int, length: int) -> dict | None: + directory = os.environ.get("BEAM_DENSE_BLOCK_SOURCE_CAPTURES") + if not directory: + return None + captures = torch.load(Path(directory) / f"seed{seed}" / f"T{length}.pt", map_location="cuda") + x, positions = case_inputs(seed, length) + assert torch.equal(captures["input"], x), "the generated input differs from the captured one" + assert torch.equal( + captures["positions"].view(1, length), positions.to(captures["positions"].dtype) + ) + return captures + + +# ----------------------------------------------------------------------------- replays + + +def replay( + blk: BeamDenseBlockBf16Sm103Tp1, md: TrtllmAttentionMetadata, src: dict, label: str +) -> list[dict]: + """Each operation of the block on `src`'s own input to it, against `src`'s output of it.""" + length = src["input"].shape[1] + hidden = BLOCK_CONFIG["hidden_size"] + heads, kv_heads = BLOCK_CONFIG["num_attention_heads"], BLOCK_CONFIG["num_key_value_heads"] + d = BLOCK_CONFIG["head_dim"] + eps = BLOCK_CONFIG["rms_norm_eps"] + w = dict(blk.checkpoint_tensors()) + positions = src["positions"].view(length).cuda() + step = blk.step_args(md) + # Logical block 0 is the manager's local slot 0: the core call indexes the slot. + assert step["local_layer_idx"] == md.kv_cache_manager.layer_offsets[LAYER_IDX] == 0 + + def flat(name: str, width: int) -> torch.Tensor: + # Dense row-major: the GEMM and norm entries read their input that way. + return src[name].reshape(length, width).contiguous() + + def check(name: str, rule: str | None, actual: torch.Tensor, expected: str) -> dict: + result = compare(name, actual.reshape(src[expected].shape), src[expected], rule) + result.update(reference=label) + return result + + x = flat("input", hidden) + h = flat("input_layernorm", hidden) + out = [ + check( + "input_layernorm", + "rmsnorm", + flashinfer_rmsnorm(x, w["input_layernorm.weight"], eps), + "input_layernorm", + ) + ] + for proj in ("q_proj", "k_proj", "v_proj", "gate_proj"): + weight = w[f"self_attn.{proj}.weight"] + out.append(check(proj, "projection", cublas_mm(h, weight.t()), proj)) + q_rows = flat("q_proj", heads * d).reshape(length * heads, d) + k_rows = flat("k_proj", kv_heads * d).reshape(length * kv_heads, d) + out.append( + check( + "q_norm", + "rmsnorm", + flashinfer_rmsnorm(q_rows, w["self_attn.q_norm.weight"], eps), + "q_norm", + ) + ) + out.append( + check( + "k_norm", + "rmsnorm", + flashinfer_rmsnorm(k_rows, w["self_attn.k_norm.weight"], eps), + "k_norm", + ) + ) + q_rope, k_rope = beam_neox_rope( + src["q_norm"].reshape(length, heads, d), + src["k_norm"].reshape(length, kv_heads, d), + positions, + ROPE_THETA, + 1.0, + ) + out.append(check("q_rope", "rope", q_rope, "q_rope")) + out.append(check("k_rope", "rope", k_rope, "k_rope")) + qkv = torch.cat( + (flat("q_rope", heads * d), flat("k_rope", kv_heads * d), flat("v_proj", kv_heads * d)), + dim=-1, + ) + core = torch.empty(length, heads * d, dtype=DTYPE, device="cuda") + core = thop_attention( + q=qkv, + output=core, + num_heads=heads, + num_kv_heads=kv_heads, + head_size=d, + attention_window_size=BLOCK_CONFIG["sliding_window"], + **step, + **_attention.CALL_CONSTANTS, + ) + # No rule of its own: the attention core is judged through the attention module below. + out.append(check("attention_core", None, core, "attention_core")) + gated = beam_output_gate( + src["attention_core"].reshape(length, heads, d), + src["gate_proj"].reshape(length, heads, d), + w["self_attn.gate_bias"], + ) + out.append(check("gated", "gate", gated, "gated")) + out.append( + check( + "o_proj", + "projection", + cublas_mm(flat("gated", heads * d), w["self_attn.o_proj.weight"].t()), + "attention_output", + ) + ) + out.append( + check( + "attention_output", + "attention_output", + blk.attention(h, positions, step), + "attention_output", + ) + ) + out.append( + check( + "post_attention_layernorm", + "rmsnorm", + flashinfer_rmsnorm( + flat("attention_output", hidden), w["post_attention_layernorm.weight"], eps + ), + "post_attention_layernorm", + ) + ) + out.append( + check( + "attention_residual", + "residual", + beam_depth_residual(flat("post_attention_layernorm", hidden), x, *blk.attention_scales), + "attention_residual", + ) + ) + out.append( + check( + "pre_feedforward_layernorm", + "rmsnorm", + flashinfer_rmsnorm( + flat("attention_residual", hidden), w["pre_feedforward_layernorm.weight"], eps + ), + "pre_feedforward_layernorm", + ) + ) + h2 = flat("pre_feedforward_layernorm", hidden) + inter = BLOCK_CONFIG["intermediate_size"] + out.append( + check( + "mlp_gate", + "projection", + cublas_mm(h2, w["mlp.gate_proj.weight"].t()), + "mlp_gate", + ) + ) + out.append(check("mlp_up", "projection", cublas_mm(h2, w["mlp.up_proj.weight"].t()), "mlp_up")) + act = flashinfer_silu_and_mul( + torch.cat((flat("mlp_gate", inter), flat("mlp_up", inter)), dim=-1) + ) + out.append(check("mlp_act", "swiglu_activation", act, "mlp_act")) + out.append( + check( + "down_proj", + "projection", + cublas_mm(flat("mlp_act", inter), w["mlp.down_proj.weight"].t()), + "mlp_output", + ) + ) + out.append(check("mlp_output", "swiglu_output", blk.feed_forward(h2), "mlp_output")) + out.append( + check( + "post_feedforward_layernorm", + "rmsnorm", + flashinfer_rmsnorm( + flat("mlp_output", hidden), w["post_feedforward_layernorm.weight"], eps + ), + "post_feedforward_layernorm", + ) + ) + out.append( + check( + "block_residual", + "residual", + beam_depth_residual( + flat("post_feedforward_layernorm", hidden), + flat("attention_residual", hidden), + *blk.feed_forward_scales, + ), + "block_output", + ) + ) + return out + + +def accounting(seed: int, blk: BeamDenseBlockBf16Sm103Tp1) -> dict: + """The fifteen tensors: the same names, order, shapes, dtypes and values on both sides.""" + state = generated(seed)["state"] + expected = reference.param_shapes() + loaded = blk.checkpoint_tensors() + assert list(loaded) == list(state) == list(expected), (list(loaded), list(state)) + for name, shape in expected.items(): + assert tuple(loaded[name].shape) == shape and loaded[name].dtype == DTYPE, name + assert torch.equal(loaded[name], state[name]), f"{name} differs from the reference's" + assert not reference.state_problems(generated(seed)["state_cpu"]) + record = {"tensors": len(loaded), "state_problems": []} + if os.environ.get("BEAM_DENSE_BLOCK_EVIDENCE"): + record["sha256"] = { + name: tensor_sha256(t) for name, t in generated(seed)["state_cpu"].items() + } + record["input_sha256"] = tensor_sha256(generated(seed)["input"]) + return record + + +# ----------------------------------------------------------------------------- op-path instrumentation + + +@contextlib.contextmanager +def recording(*wrappers): + """Log every call of these entries, in order, with the arguments that identify the path.""" + calls: list[dict] = [] + saved = [] + for w in wrappers: + cls = type(w) + original = cls.__call__ + saved.append((cls, original)) + + def logged(self, *args, _original=original, _name=cls.__name__, **kwargs): + entry = {"op": _name} + tensors = [a for a in args if isinstance(a, torch.Tensor)] + entry["shapes"] = [list(t.shape) for t in tensors] + for key in ( + "num_heads", + "num_kv_heads", + "head_size", + "attention_window_size", + "max_seq_len", + "tokens_per_block", + "mask_type", + "is_fused_qkv", + "num_contexts", + "num_ctx_tokens", + "use_paged_context_fmha", + "position_embedding_type", + "local_layer_idx", + "quant_mode", + ): + if key in kwargs: + entry[key] = kwargs[key] + if "q" in kwargs: + entry["q_shape"] = list(kwargs["q"].shape) + entry["qkv"] = kwargs["q"] + calls.append(entry) + return _original(self, *args, **kwargs) + + cls.__call__ = logged + try: + yield calls + finally: + for cls, original in reversed(saved): + cls.__call__ = original + + +@contextlib.contextmanager +def builtin_beam_forbidden(): + """The built-in Beam modeling must not run while the component does.""" + targets = [ + builtin_beam.BeamDecoderLayer, + builtin_beam.BeamAttention, + builtin_beam.BeamSwiGLU, + builtin_beam.BeamRMSNorm, + ] + saved = [(cls, cls.forward) for cls in targets] + + def refuse(self, *args, **kwargs): + raise AssertionError(f"the built-in {type(self).__name__} ran inside the component") + + for cls in targets: + cls.forward = refuse + try: + yield + finally: + for cls, forward in saved: + cls.forward = forward + + +@contextlib.contextmanager +def no_dispatch(): + """Fail on any op call: rejected inputs must be rejected before a kernel runs.""" + saved = [] + for w in OPS: + cls = type(w) + saved.append((cls, cls.__call__)) + + def refuse(self, *args, _name=cls.__name__, **kwargs): + raise AssertionError(f"{_name} dispatched before the input was rejected") + + cls.__call__ = refuse + try: + yield + finally: + for cls, original in reversed(saved): + cls.__call__ = original + + +def cache_rows(md: TrtllmAttentionMetadata, length: int) -> tuple[torch.Tensor, torch.Tensor]: + """This request's K and V rows read back from the pool, [T, Hkv, D] each.""" + mgr = md.kv_cache_manager + view = mgr.get_buffers(LAYER_IDX, kv_layout="HND") + view = view[:: mgr.get_layer_page_index_scale(LAYER_IDX) // 2] + tpb = KV_CACHE["tokens_per_block"] + keys, values = [], [] + for page in range((length + tpb - 1) // tpb): + k_off = int(md.kv_cache_block_offsets[0, 0, 0, page]) + rows = min(tpb, length - page * tpb) + keys.append(view[k_off // 2, 0, :, :rows, :].permute(1, 0, 2)) + values.append(view[k_off // 2, 1, :, :rows, :].permute(1, 0, 2)) + return torch.cat(keys), torch.cat(values) + + +# ----------------------------------------------------------------------------- tests: configuration and construction + + +def test_configuration_is_the_fixed_scope() -> None: + """One configuration, no open choices: what the component accepts is what is tested.""" + assert CASES == ((17, 1), (17, 17), (17, 129), (41, 1), (41, 17), (41, 129)) + assert block_module.SUPPORTED_LAYER_IDX == LAYER_IDX + for field, value in block_module.SUPPORTED_CONFIG.items(): + assert BLOCK_CONFIG[field] == value, field + assert BLOCK_CONFIG["rope_parameters"]["sliding_attention"] == block_module.SUPPORTED_ROPE + assert block_module.MAX_TOKENS == 1024 >= max(LENGTHS) + assert reference.BLOCK0.hidden == BLOCK_CONFIG["hidden_size"] + assert reference.BLOCK0.intermediate == BLOCK_CONFIG["intermediate_size"] + assert (reference.BLOCK0.heads, reference.BLOCK0.kv_heads, reference.BLOCK0.head_dim) == ( + 48, + 8, + 128, + ) + assert reference.BLOCK0.window == BLOCK_CONFIG["sliding_window"] == KV_CACHE["max_seq_len"] + assert (reference.SEEDS, reference.LENGTHS) == (SEEDS, LENGTHS) + assert torch.cuda.get_device_capability() == SM, "this suite certifies sm_103 only" + + +def test_block_constructs_alone(ieee_fp32) -> None: + """The block is built directly from its configuration, without the model around it.""" + + def refuse(self, *args, **kwargs): + raise AssertionError(f"{type(self).__name__} constructed") + + with pytest.MonkeyPatch.context() as mp: + for cls in ( + builtin_beam.BeamForCausalLM, + builtin_beam.BeamModel, + builtin_beam.BeamDecoderLayer, + ): + mp.setattr(cls, "__init__", refuse) + blk = BeamDenseBlockBf16Sm103Tp1(block_config(), LAYER_IDX, device="cuda") + assert isinstance(blk, torch.nn.Module) and not list(blk.children()) + assert list(blk.checkpoint_tensors()) == list(reference.param_shapes()) + for name, shape in reference.param_shapes().items(): + t = blk.checkpoint_tensors()[name] + assert tuple(t.shape) == shape and t.dtype == DTYPE and t.is_cuda, name + q_width = BLOCK_CONFIG["num_attention_heads"] * BLOCK_CONFIG["head_dim"] + kv_width = BLOCK_CONFIG["num_key_value_heads"] * BLOCK_CONFIG["head_dim"] + assert ( + blk.checkpoint_tensors()["self_attn.q_proj.weight"].shape == (q_width, 5120) == (6144, 5120) + ) + assert blk.checkpoint_tensors()["self_attn.gate_proj.weight"].shape == (6144, 5120) + assert ( + blk.checkpoint_tensors()["self_attn.k_proj.weight"].shape + == (kv_width, 5120) + == (1024, 5120) + ) + assert blk.checkpoint_tensors()["self_attn.v_proj.weight"].shape == (1024, 5120) + assert blk.attention_scales == (1 / math.sqrt(2), 1 / math.sqrt(2)) + assert blk.feed_forward_scales == (1 / math.sqrt(3), math.sqrt(2) / math.sqrt(3)) + assert "modeling_beam" not in Path(block_module.__file__).read_text() + + +# ----------------------------------------------------------------------------- tests: declared limits + +_BAD_CONFIGS = [ + ("hidden_size", dict(hidden_size=4096)), + ("intermediate_size", dict(intermediate_size=18432)), + ("num_attention_heads", dict(num_attention_heads=32)), + ("num_key_value_heads", dict(num_key_value_heads=4)), + ("head_dim", dict(head_dim=64)), + ("rms_norm_eps", dict(rms_norm_eps=1e-5)), + ("hidden_act", dict(hidden_act="gelu")), + ("attention_bias", dict(attention_bias=True)), + ("sliding_window", dict(sliding_window=1024)), + ("sliding_attention layer", dict(layer_types=["full_attention"])), + ("dense MLP", dict(mlp_layer_types=["sparse"])), + ( + "rope parameters", + dict(rope_parameters={"sliding_attention": {"rope_type": "default", "rope_theta": 1e6}}), + ), + ( + "rope parameters", + dict( + rope_parameters={ + "sliding_attention": {"rope_type": "linear", "rope_theta": 10000.0, "factor": 2.0} + } + ), + ), +] + + +@pytest.mark.parametrize( + "match,overrides", _BAD_CONFIGS, ids=[f"{m}-{i}" for i, (m, _) in enumerate(_BAD_CONFIGS)] +) +def test_rejects_unsupported_configuration(match, overrides) -> None: + with pytest.raises(ValueError, match=match): + BeamDenseBlockBf16Sm103Tp1(block_config(**overrides), LAYER_IDX, device="cuda") + + +def test_rejects_other_blocks_and_devices() -> None: + for layer_idx in (3, 5): + with pytest.raises(ValueError, match="only dense blocks"): + BeamDenseBlockBf16Sm103Tp1(block_config(), layer_idx, device="cuda") + with pytest.raises(ValueError, match="sm_103"): + BeamDenseBlockBf16Sm103Tp1(block_config(), LAYER_IDX, device="cpu") + + +def test_weights_must_be_exactly_the_blocks() -> None: + blk = BeamDenseBlockBf16Sm103Tp1(block_config(), LAYER_IDX, device="cuda") + good = { + name: torch.zeros(shape, dtype=DTYPE) for name, shape in reference.param_shapes().items() + } + missing = dict(good) + del missing["self_attn.gate_bias"] + extra = dict(good, **{"mlp.experts.0.down_proj.weight": torch.zeros(4, dtype=DTYPE)}) + transposed = dict(good, **{"self_attn.q_proj.weight": torch.zeros(5120, 6144, dtype=DTYPE)}) + fp32 = dict(good, **{"self_attn.q_norm.weight": torch.zeros(128, dtype=torch.float32)}) + for weights, match in ( + (missing, "missing"), + (extra, "unexpected"), + (transposed, "q_proj"), + (fp32, "q_norm"), + ): + with pytest.raises(ValueError, match=match): + blk.load_weights(weights) + + +_BAD_INPUTS = [ + "fp16_hidden", + "fp32_hidden", + "hidden_5120x2", + "batch_2", + "too_long", + "empty", + "non_contiguous", + "positions_offset", + "float_positions", + "positions_shape", + "metadata_type", + "cached_prefix", + "two_requests", + "context_length_mismatch", + "kv_cache_heads", + "kv_cache_dtype", +] + + +@pytest.mark.parametrize("case", _BAD_INPUTS) +def test_rejects_unsupported_inputs_before_dispatch(case) -> None: + """Each declared restriction raises ValueError, and no op runs first.""" + blk = BeamDenseBlockBf16Sm103Tp1(block_config(), LAYER_IDX, device="cuda") + length = 17 + x = torch.randn(1, length, 5120, device="cuda").to(DTYPE) + positions = torch.arange(length, device="cuda").view(1, length) + context = dict(num_tokens=length) + if case == "fp16_hidden": + x = x.half() + elif case == "fp32_hidden": + x = x.float() + elif case == "hidden_5120x2": + x = torch.cat((x, x), dim=-1) + elif case == "batch_2": + x = torch.cat((x, x), dim=0) + elif case == "too_long": + length = block_module.MAX_TOKENS + 1 + x = torch.randn(1, length, 5120, device="cuda").to(DTYPE) + positions = torch.arange(length, device="cuda").view(1, length) + context = dict(num_tokens=17) + elif case == "empty": + x = x[:, :0] + positions = positions[:, :0] + elif case == "non_contiguous": + x = torch.randn(1, length, 2 * 5120, device="cuda").to(DTYPE)[..., ::2] + elif case == "positions_offset": + positions = positions + 1 + elif case == "float_positions": + positions = positions.float() + elif case == "positions_shape": + positions = positions.view(length) + elif case == "cached_prefix": + context = dict(num_tokens=length, cached_tokens=8) + elif case == "two_requests": + context = dict(num_tokens=length, num_requests=2) + elif case == "context_length_mismatch": + context = dict(num_tokens=length - 1) + elif case == "kv_cache_heads": + context = dict(num_tokens=length, num_kv_heads=4) + elif case == "kv_cache_dtype": + context = dict(num_tokens=length, dtype=DataType.HALF) + with fresh_context(**context) as md: + if case == "metadata_type": + md = object() + with no_dispatch(), pytest.raises(ValueError): + blk(x, positions, md) + + +# ----------------------------------------------------------------------------- tests: the six cases + + +@pytest.mark.parametrize("seed,length", CASES, ids=CASE_IDS) +def test_block_matches_reference(seed, length, ieee_fp32) -> None: + """One case: the whole block, then every operation on its own input.""" + blk = loaded_block(seed) + accounting_record = accounting(seed, blk) + ref = reference_record(seed, length) + x, positions = case_inputs(seed, length) + results = [] + with fresh_context(length) as md: + out = blk(x, positions, md) + torch.cuda.synchronize() + assert out.shape == (1, length, 5120) and out.dtype == DTYPE + final = compare("block_output", out, ref["block_output"], "block_output") + final.update(reference="reference") + results.append(final) + results.extend(replay(blk, md, ref, "reference")) + captures = source_captures(seed, length) + if captures is not None: + from_source = compare("block_output", out, captures["block_output"], "block_output") + from_source.update(reference="source_captures") + results.append(from_source) + results.extend(replay(blk, md, captures, "source_captures")) + _write_evidence( + f"block-seed{seed}-T{length}", + { + "test": "test_block_matches_reference", + "case": {"seed": seed, "length": length}, + "accounting": accounting_record, + "rules": {k: list(v) for k, v in RULES.items()}, + "results": results, + "environment": _environment(ieee_fp32), + }, + ) + assert_all_pass(results) + + +@pytest.mark.parametrize("seed,length", CASES, ids=CASE_IDS) +def test_forward_runs_the_component_path(seed, length, ieee_fp32) -> None: + """One case: the component's own composition ran, op by op, through thop attention on + this GPU, with no SDPA and no built-in Beam module on the way.""" + blk = loaded_block(seed) + x, positions = case_inputs(seed, length) + activities = [torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] + with fresh_context(length) as md: + with builtin_beam_forbidden(), validating(*OPS), recording(*OPS) as calls: + with torch.profiler.profile(activities=activities) as prof: + out = blk(x, positions, md) + torch.cuda.synchronize() + keys, values = cache_rows(md, length) + mgr = md.kv_cache_manager + cache_mapping = { + "layer_offsets": dict(mgr.layer_offsets), + "num_local_layers": mgr.num_local_layers, + "num_layers": mgr.num_layers, + } + events = prof.events() + kernels = [e.name for e in events if e.device_type == torch.autograd.DeviceType.CUDA] + cpu_ops = sorted({e.name for e in events if e.device_type == torch.autograd.DeviceType.CPU}) + names = [c["op"] for c in calls] + attention = [c for c in calls if c["op"] == "ThopAttention"] + qkv = attention[0].pop("qkv") if attention else None + for c in calls: + c.pop("qkv", None) + record = { + "test": "test_forward_runs_the_component_path", + "case": {"seed": seed, "length": length}, + "block": f"{type(blk).__module__}.{type(blk).__qualname__}", + "op_calls": calls, + "cuda_kernels": kernels, + "attention_kernels": [k for k in kernels if _ATTENTION_KERNEL.search(k)], + "sdpa_cpu_ops": [n for n in cpu_ops if _SDPA_OPS.search(n)], + "kv_cache": dict(KV_CACHE, manager=type(md.kv_cache_manager).__name__, **cache_mapping), + "environment": _environment(ieee_fp32), + } + _write_evidence(f"path-seed{seed}-T{length}", record) + + assert type(blk).__module__.endswith("dense_block_bf16__sm_103__tp1.block") + assert tuple(names) == EXPECTED_CALLS, names + (call,) = attention + assert call["q_shape"] == [length, (48 + 2 * 8) * 128] + assert { + k: call[k] for k in ("num_heads", "num_kv_heads", "head_size", "attention_window_size") + } == { + "num_heads": 48, + "num_kv_heads": 8, + "head_size": 128, + "attention_window_size": 512, + } + assert call["mask_type"] == 1 and call["is_fused_qkv"] and call["position_embedding_type"] == 0 + assert call["num_contexts"] == 1 and call["num_ctx_tokens"] == length + assert call["use_paged_context_fmha"] is False and call["quant_mode"] == 0 + assert isinstance(md.kv_cache_manager, KVCacheManagerV2) + # Logical block 0 in a one-layer manager: slot 0, which the core call indexes. + assert cache_mapping == {"layer_offsets": {0: 0}, "num_local_layers": 1, "num_layers": 1} + assert call["local_layer_idx"] == 0 + assert record["attention_kernels"], "no attention kernel ran on the GPU" + assert not record["sdpa_cpu_ops"], record["sdpa_cpu_ops"] + assert torch.cuda.get_device_capability(x.device) == SM + assert out.shape == (1, length, 5120) and bool(torch.isfinite(out.float()).all()) + # The op appended exactly this call's post-RoPE K and V to the pool. + k_in = qkv[:, 48 * 128 : 56 * 128].reshape(length, 8, 128) + v_in = qkv[:, 56 * 128 :].reshape(length, 8, 128) + assert torch.equal(keys, k_in) and torch.equal(values, v_in) + + +def test_detects_a_wrong_residual_coefficient(ieee_fp32, monkeypatch) -> None: + """The complete-block comparison rejects a zeroed final residual coefficient and + accepts the unchanged block, on the same case and the same comparison.""" + seed, length = 17, 17 + blk = loaded_block(seed) + ref = reference_record(seed, length) + x, positions = case_inputs(seed, length) + runs = {} + for label in ("unchanged", "faulted", "restored"): + if label == "faulted": + branch_scale, _ = blk.feed_forward_scales + monkeypatch.setattr(blk, "feed_forward_scales", (branch_scale, 0.0)) + elif label == "restored": + monkeypatch.undo() + with fresh_context(length) as md: + out = blk(x, positions, md) + runs[label] = compare("block_output", out, ref["block_output"], "block_output") + runs[label]["feed_forward_scales"] = list(blk.feed_forward_scales) + _write_evidence( + f"fault-seed{seed}-T{length}", + { + "test": "test_detects_a_wrong_residual_coefficient", + "case": {"seed": seed, "length": length}, + "fault": "final residual coefficient sqrt(2)/sqrt(3) -> 0.0, fixture-scoped", + "runs": runs, + "environment": _environment(ieee_fp32), + }, + ) + assert runs["unchanged"]["passed"] is True, runs["unchanged"]["problems"] + assert runs["faulted"]["passed"] is False and runs["faulted"]["violations"] > 0 + assert runs["restored"][ + "passed" + ] is True and blk.feed_forward_scales == block_module.depth_scales(0, False) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_batching.py b/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_batching.py new file mode 100644 index 000000000000..83fdec764682 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_batching.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Batched engine steps of the Beam attention contract: request spans, rejections and rows. + +The native runtime orders a step's requests contexts first, then generations; each request +has its own cached prefix and KV length. ``check_batch_metadata`` must return each request's +``(past, new tokens)`` from the prepared metadata and reject a step whose rows or requests +disagree; ``check_metadata`` keeps the one-request contract; the full model selects the last +row of each request for the sampler. Needs one CUDA device (a small KVCacheManagerV2); no +model files. +""" + +from __future__ import annotations + +import contextlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam import _attention +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.model import ( # noqa: E501 + BeamFullModelNvfp4Sm103Tp4, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping + +KV_HEADS, HEAD_DIM, TOKENS_PER_BLOCK, MAX_SEQ_LEN = 2, 128, 32, 4096 +#: (cached tokens, new tokens) per request: two contexts (fresh, after a cached prefix) and +#: two generation requests at different depths. +CONTEXTS = ((0, 40), (64, 20)) +GENERATIONS = ((100, 1), (7, 1)) + + +@contextlib.contextmanager +def prepared_step(contexts, generations, *, cache_reuse=True): + """A one-layer BF16 KVCacheManagerV2 and prepared metadata for one engine step.""" + if not torch.cuda.is_available(): + pytest.skip("needs a CUDA device") + requests = list(contexts) + list(generations) + mgr = KVCacheManagerV2( + KvCacheConfig(max_tokens=64 * TOKENS_PER_BLOCK * len(requests), enable_block_reuse=False), + CacheTypeCpp.SELF, + num_layers=1, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=len(requests), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.BF16, + ) + try: + request_ids = list(range(1, len(requests) + 1)) + mgr.add_dummy_requests(request_ids, token_nums=[past + n for past, n in requests]) + md = TrtllmAttentionMetadata( + max_num_requests=len(requests), + max_num_tokens=8192, + kv_cache_manager=mgr, + num_contexts=len(contexts), + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=[past for past, _ in requests] + ), + seq_lens=torch.tensor([n for _, n in requests], dtype=torch.int), + request_ids=request_ids, + # The engine's context length of a context request is its new tokens (this chunk). + prompt_lens=[n for _, n in contexts] + [past + n for past, n in generations], + runtime_features=AttentionRuntimeFeatures(cache_reuse=cache_reuse), + ) + md.prepare() + yield md + finally: + mgr.shutdown() + + +def _spans(md, num_tokens): + return _attention.check_batch_metadata( + md, num_tokens, layer_idx=0, num_kv_heads=KV_HEADS, head_dim=HEAD_DIM + ) + + +def test_mixed_step_spans(): + requests = list(CONTEXTS) + list(GENERATIONS) + with prepared_step(CONTEXTS, GENERATIONS) as md: + assert _spans(md, sum(n for _, n in requests)) == requests + + +def test_generation_only_step_spans(): + with prepared_step((), GENERATIONS) as md: + assert _spans(md, len(GENERATIONS)) == list(GENERATIONS) + + +def test_one_request_contract_is_kept(): + with prepared_step(CONTEXTS[1:], ()) as md: + past = _attention.check_metadata( + md, CONTEXTS[1][1], layer_idx=0, num_kv_heads=KV_HEADS, head_dim=HEAD_DIM + ) + assert past == CONTEXTS[1][0] + with prepared_step(CONTEXTS, GENERATIONS) as md: + with pytest.raises(ValueError, match="one request per call"): + _attention.check_metadata(md, 62, layer_idx=0, num_kv_heads=KV_HEADS, head_dim=HEAD_DIM) + + +@pytest.mark.parametrize("num_tokens", [61, 63, 1]) +def test_rows_must_be_the_requests_new_tokens(num_tokens): + with prepared_step(CONTEXTS, GENERATIONS) as md: + with pytest.raises(ValueError, match="do not fill"): + _spans(md, num_tokens) + + +def test_cached_prefix_needs_paged_context_fmha(): + with prepared_step(CONTEXTS, GENERATIONS, cache_reuse=False) as md: + with pytest.raises(ValueError, match="paged-context FMHA"): + _spans(md, 62) + + +def test_generation_request_appends_one_token(): + with prepared_step(CONTEXTS, ((100, 2),)) as md: + with pytest.raises(ValueError, match="exactly one token"): + _spans(md, 62) + + +def test_unheld_layer_is_rejected(): + with prepared_step(CONTEXTS, GENERATIONS) as md: + with pytest.raises(ValueError, match="does not hold logical layer 3"): + _attention.check_batch_metadata( + md, 62, layer_idx=3, num_kv_heads=KV_HEADS, head_dim=HEAD_DIM + ) + + +@pytest.mark.parametrize( + "spans,expected", + [ + ([(0, 40)], [39]), + ([(0, 40), (64, 20), (100, 1), (7, 1)], [39, 59, 60, 61]), + ([(5, 1), (9, 1), (2, 1)], [0, 1, 2]), + ], +) +def test_default_rows_are_each_requests_last_row(spans, expected): + select = BeamFullModelNvfp4Sm103Tp4._select_rows + assert select(None, None, spans) == expected + assert select(None, "all", spans) == list(range(sum(n for _, n in spans))) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_dflash2.py b/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_dflash2.py new file mode 100644 index 000000000000..401a5290a719 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_dflash2.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CUDA tests of the pieces the Beam NVFP4 target composes for DFlash 2 speculative decoding. + +Generated inputs, one CUDA device, no model files. The pieces are the native ones the target +reuses -- the DFlash worker's greedy verifier and candidate merge, the DFlash 2 selector, the +one-model speculative sampler, the DFlash capture buffer -- plus the drafter's logit softcap +the target installs (``DFlashSoftcapLogitsProcessor``). What is checked: + +* the drafter's unary logits are ``cap * tanh(raw / cap)`` exactly once, and the candidate + path the selector picks discriminates an omitted and a doubled cap (two cases whose + bilinear bonus sits between the once-capped and the faulty scores), with the candidates on + different TP vocabulary shards merged by the worker's global top-k; +* greedy verification on identical logits: zero, partial and full acceptance of seven + proposals in one four-request step with unequal counts, the target's argmax as correction + or bonus token, ties to the lowest id, and the committed length ``1 + accepted``; +* stops inside a verification block: an end id among the accepted tokens or as the bonus + token, and completion budgets 1, 7, 8, 9 and 64 against a full block, emit nothing past + the stop or budget; +* the capture buffer: each capture layer's rows land in its own 5120-wide slice in ascending + layer order, and a capture at a neighbouring layer is detected. +""" + +from __future__ import annotations + +import types + +import pytest +import torch + +from tensorrt_llm._torch.distributed import ops as dist_ops +from tensorrt_llm._torch.models.modeling_dflash import ( + DFlash2CandidateSelector, + DFlashForCausalLM, + DFlashSoftcapLogitsProcessor, + dflash_final_logit_softcap, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState +from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata, DFlashWorker +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler +from tensorrt_llm.bindings.executor import FinishReason +from tensorrt_llm.bindings.executor import SamplingConfig as ExecutorSamplingConfig +from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + +K = 7 # proposals per block: the anchor plus seven rows of a verification block +CAP = 30.0 + + +def _worker(tp_rank: int = 0, tp_size: int = 1) -> DFlashWorker: + config = DFlashDecodingConfig(max_draft_len=K, attention_backend="TRTLLM") + return DFlashWorker(config, Mapping(world_size=tp_size, tp_size=tp_size, rank=tp_rank)) + + +def _spec_metadata(num_requests: int = 8) -> DFlashSpecMetadata: + md = DFlashSpecMetadata( + max_draft_len=K, + max_total_draft_tokens=K, + spec_dec_mode=SpeculativeDecodingMode.DFLASH, + max_num_requests=num_requests, + ) + md.runtime_draft_len = K + md.runtime_tokens_per_gen_step = K + 1 + return md + + +# --------------------------------------------------------------------------------------------- +# Draft logit softcap and the DFlash 2 candidate selector +# --------------------------------------------------------------------------------------------- + + +@torch.no_grad() +def test_softcap_processor_caps_the_head_output_once(): + head = torch.nn.Linear(16, 24, bias=False, device="cuda", dtype=torch.bfloat16) + hidden = 8 * torch.randn(5, 16, device="cuda", dtype=torch.bfloat16) + raw = head(hidden).float() + capped = DFlashSoftcapLogitsProcessor(CAP)(hidden, head, None, True) + assert capped.dtype == torch.float32 + assert torch.equal(capped, CAP * torch.tanh(raw / CAP)) + assert float(capped.abs().max()) < CAP + assert torch.equal(LogitsProcessor()(hidden, head, None, True), raw) + + +def test_softcap_value_comes_from_the_drafter_config(): + config = types.SimpleNamespace(dflash_config={"final_logit_softcapping": 30}) + assert dflash_final_logit_softcap(config) == 30.0 + assert dflash_final_logit_softcap(types.SimpleNamespace(dflash_config={})) is None + for bad in (0, -1.0, "30", True): + with pytest.raises(ValueError): + dflash_final_logit_softcap( + types.SimpleNamespace(dflash_config={"final_logit_softcapping": bad}) + ) + with pytest.raises(ValueError): + DFlashSoftcapLogitsProcessor(0.0) + + +VOCAB, SHARDS = 32, 4 +SHARD = VOCAB // SHARDS +FIRST, SECOND, ANCHOR = 5, 27, 12 # candidates on shards 0 and 3; the anchor on shard 1 +HIDDEN = 16 + + +def _selector(bonus: float) -> DFlash2CandidateSelector: + """A rank-2 selector whose bilinear term adds ``bonus`` to SECOND after the anchor only.""" + predecessor = torch.zeros(VOCAB, 2, dtype=torch.bfloat16, device="cuda") + predecessor[ANCHOR, 0] = 1.0 + successor = torch.zeros(VOCAB, 2, dtype=torch.bfloat16, device="cuda") + successor[SECOND, 0] = bonus + projection = torch.zeros(2, HIDDEN, dtype=torch.bfloat16, device="cuda") + projection[0, 0] = 1.0 + return DFlash2CandidateSelector(predecessor, successor, projection, 2, VOCAB, 2) + + +class _ShardHead(torch.nn.Module): + """Stands for one TP rank's vocabulary-parallel draft head: returns that rank's raw slice.""" + + def __init__(self, raw: torch.Tensor, rank: int): + super().__init__() + self.raw = raw[..., rank * SHARD : (rank + 1) * SHARD] + + def forward(self, hidden): + return self.raw.expand(hidden.shape[0], SHARD) + + +def _selected_token(processor_layers: int, bonus: float, monkeypatch) -> int: + """The draft token the native path picks for one block position, emulating four TP ranks. + + Each rank applies the draft logits processor ``processor_layers`` times (0 = the uncapped + native processor, 1 = the target's capped one, 2 = a doubled cap) to its vocabulary shard; + the worker's global top-k merges the shards through an all-gather of every rank's + (id, value) pairs; the selector scores the merged candidates after the anchor; the worker's + greedy draft sampler reads the result. + """ + raw = torch.full((VOCAB,), -100.0, device="cuda") + raw[FIRST], raw[SECOND] = 60.0, 20.0 + hidden = torch.zeros(1, HIDDEN, dtype=torch.bfloat16, device="cuda") + hidden[0, 0] = 1.0 + spec = types.SimpleNamespace(draft_vocab_size=VOCAB, vocab_size=VOCAB) + + def shard_logits(rank): + processor = ( + LogitsProcessor() if processor_layers == 0 else DFlashSoftcapLogitsProcessor(CAP) + ) + logits = processor(hidden, _ShardHead(raw, rank), None, True) + if processor_layers == 2: + logits = CAP * torch.tanh(logits / CAP) + return logits.view(1, 1, SHARD) + + gathered = [] + monkeypatch.setattr(dist_ops, "allgather", lambda t, mapping, dim=-1: gathered.append(t) or t) + for rank in range(SHARDS): + _worker(rank, SHARDS)._dflash2_global_top_k(shard_logits(rank), spec, 2, VOCAB) + merged = torch.cat(gathered, dim=-1) + monkeypatch.setattr(dist_ops, "allgather", lambda t, mapping, dim=-1: merged) + picks = set() + for rank in range(SHARDS): + worker = _worker(rank, SHARDS) + candidate_ids, unary, block_logits = worker._dflash2_global_top_k( + shard_logits(rank), spec, 2, VOCAB + ) + assert sorted(candidate_ids.view(-1).tolist()) == [FIRST, SECOND] + selector = types.SimpleNamespace(candidate_selector=_selector(bonus)) + block_logits = DFlashForCausalLM.select_candidate_path( + selector, + block_logits, + candidate_ids, + unary, + hidden.view(1, 1, HIDDEN), + torch.tensor([ANCHOR], device="cuda"), + ) + draft_spec = types.SimpleNamespace( + wants_advanced_draft_sampling=False, draft_vocab_size=VOCAB, vocab_size=VOCAB + ) + picks.add(int(worker.sample_draft_tokens(block_logits, draft_spec, 1).view(-1)[0])) + assert len(picks) == 1, "the TP ranks disagree on the draft token" + return picks.pop() + + +def test_selector_path_detects_an_omitted_cap(monkeypatch): + """Bonus 20: capped [28.92, 17.48 + 20] picks SECOND; uncapped [60, 20 + 20] picks FIRST.""" + assert _selected_token(1, 20.0, monkeypatch) == SECOND + assert _selected_token(0, 20.0, monkeypatch) == FIRST + assert _selected_token(2, 20.0, monkeypatch) == SECOND + + +def test_selector_path_detects_a_doubled_cap(monkeypatch): + """Bonus 9: capped [28.92, 17.48 + 9] picks FIRST; doubled [22.38, 15.74 + 9] picks SECOND.""" + assert _selected_token(1, 9.0, monkeypatch) == FIRST + assert _selected_token(2, 9.0, monkeypatch) == SECOND + assert _selected_token(0, 9.0, monkeypatch) == FIRST + + +# --------------------------------------------------------------------------------------------- +# Greedy verification on identical logits +# --------------------------------------------------------------------------------------------- + +V = 64 + + +def _rows(argmaxes) -> torch.Tensor: + """FP32 logits rows whose argmax is each of ``argmaxes``.""" + logits = torch.linspace(-1.0, 1.0, V, device="cuda").repeat(len(argmaxes), 1) + for row, token in enumerate(argmaxes): + logits[row, token] = 50.0 + return logits + + +def _verify(contexts, blocks): + """Run the worker's acceptance on a step of ``contexts`` (argmax of each context's last + row) and ``blocks`` (per generation request: (row argmaxes [8], proposals [7])).""" + worker = _worker() + md = _spec_metadata() + rows = list(contexts) + [t for argmaxes, _ in blocks for t in argmaxes] + md.draft_tokens = torch.tensor( + [d for _, proposals in blocks for d in proposals], dtype=torch.int32, device="cuda" + ) + attn = types.SimpleNamespace(num_seqs=len(contexts) + len(blocks), num_contexts=len(contexts)) + accepted, counts = worker.sample_and_accept_draft_tokens(_rows(rows), attn, md) + return accepted.cpu(), counts.cpu() + + +def test_zero_partial_and_full_acceptance_in_one_step(): + """Four requests with 0, 3, 7 and 5 accepted proposals beside one context request.""" + target = [list(range(10 + 8 * r, 18 + 8 * r)) for r in range(4)] + proposals = [ + [63] + target[0][1:7], # first proposal wrong: nothing accepted + target[1][:3] + [63, 63, 63, 63], # three accepted, then the correction + target[2][:7], # all seven accepted, then the bonus + target[3][:5] + [63, target[3][6]], # a later match after a mismatch does not count + ] + accepted, counts = _verify([40], list(zip(target, proposals))) + assert counts.tolist() == [1, 1, 4, 8, 6] + assert accepted[0, 0] == 40 + for r, n in zip(range(4), (1, 4, 8, 6)): + # The emitted tokens are the target's own argmaxes: accepted proposals, then the + # correction (or, after full acceptance, the bonus), nothing past them. + assert accepted[1 + r, :n].tolist() == target[r][:n] + + +def test_ties_take_the_lowest_token_id(): + logits = _rows([7, 7, 7, 7, 7, 7, 7, 7]) + logits[0, 3] = 50.0 # equal maximum at a lower id than 7 + worker, md = _worker(), _spec_metadata() + md.draft_tokens = torch.tensor([7] * K, dtype=torch.int32, device="cuda") + attn = types.SimpleNamespace(num_seqs=1, num_contexts=0) + accepted, counts = worker.sample_and_accept_draft_tokens(logits, attn, md) + assert counts.tolist() == [1] and accepted[0, 0] == 3 + + +# --------------------------------------------------------------------------------------------- +# Stops inside a verification block (the one-model speculative sampler) +# --------------------------------------------------------------------------------------------- + +EOS = 60 + + +def _sampler() -> SpecSampler: + return SpecSampler( + TorchSampler.Args( + max_seq_len=8192, + max_draft_len=K, + max_num_sequences=8, + max_beam_width=1, + max_total_draft_tokens=K, + disable_overlap_scheduler=True, + ) + ) + + +class _Request(LlmRequest): + """An ``LlmRequest`` that keeps the finish reason the sampler records on it.""" + + py_finish_reason = None + + def finish_by(self, reason: FinishReason, beam: int) -> None: + super().finish_by(reason, beam) + self.py_finish_reason = reason + + +def _request(slot: int, budget: int) -> LlmRequest: + request = _Request( + request_id=slot + 1, + max_new_tokens=budget, + input_tokens=[1, 2, 3], + sampling_config=ExecutorSamplingConfig(beam_width=1, top_k=1, temperature=0.0), + is_streaming=False, + end_id=EOS, + ) + request.py_seq_slot = slot + return request + + +def _step(sampler, requests, outputs, *, is_context): + batch = ScheduledRequests() + if is_context: + batch.context_requests_last_chunk = list(requests) + else: + batch.generation_requests = list(requests) + state = sampler.sample_async(batch, outputs, [0]) + if is_context: + for request in requests: + request.context_chunk_size = request.prompt_len + request.move_to_next_context_chunk() + request.state = LlmRequestState.GENERATION_IN_PROGRESS + sampler.update_requests(state) + + +def _outputs(worker, md, accepted, counts): + """The worker's result dictionary for a step whose drafting produced no proposals.""" + batch = accepted.shape[0] + next_draft = torch.zeros(batch, K, dtype=torch.int32, device="cuda") + return { + "new_tokens": accepted, + "new_tokens_lens": counts, + "next_draft_tokens": next_draft, + "next_new_tokens": worker._prepare_next_new_tokens( + accepted, next_draft, md.batch_indices_cuda, batch, counts + ), + } + + +def _step_metadata(num_requests: int) -> DFlashSpecMetadata: + md = _spec_metadata() + md.request_ids = list(range(num_requests)) + md.prepare() # batch_indices_cuda = arange(num_requests) + return md + + +def _run_block(budgets, block_argmaxes, proposals): + """A context step (first token 9 for every request), then one verification step.""" + sampler, worker = _sampler(), _worker() + requests = [_request(i, b) for i, b in enumerate(budgets)] + md = _step_metadata(len(requests)) + attn = types.SimpleNamespace(num_seqs=len(requests), num_contexts=len(requests)) + accepted, counts = worker.sample_and_accept_draft_tokens(_rows([9] * len(requests)), attn, md) + _step(sampler, requests, _outputs(worker, md, accepted, counts), is_context=True) + live = [r for r in requests if r.state != LlmRequestState.GENERATION_COMPLETE] + if live: + md = _step_metadata(len(live)) + md.draft_tokens = torch.tensor( + [d for r in live for d in proposals[r.py_seq_slot]], dtype=torch.int32, device="cuda" + ) + attn = types.SimpleNamespace(num_seqs=len(live), num_contexts=0) + rows = [t for r in live for t in block_argmaxes[r.py_seq_slot]] + accepted, counts = worker.sample_and_accept_draft_tokens(_rows(rows), attn, md) + _step(sampler, live, _outputs(worker, md, accepted, counts), is_context=False) + return requests + + +def _generated(request): + return request.get_tokens(0)[request.prompt_len :] + + +def test_end_id_inside_and_at_the_end_of_a_block(): + block = list(range(20, 28)) + inside = block[:2] + [EOS] + block[3:] # the end id is the third accepted token + bonus = block[:7] + [EOS] # all seven proposals accepted, the bonus is the end id + requests = _run_block( + [64, 64, 64], + [inside, bonus, block], + [inside[:7], bonus[:7], block[:4] + [63, 63, 63]], + ) + assert _generated(requests[0]) == [9, 20, 21, EOS] + assert requests[0].py_finish_reason == FinishReason.END_ID + assert _generated(requests[1]) == [9] + block[:7] + [EOS] + assert requests[1].py_finish_reason == FinishReason.END_ID + # Four accepted then the correction: five tokens, still running. + assert _generated(requests[2]) == [9] + block[:5] + assert requests[2].state != LlmRequestState.GENERATION_COMPLETE + assert requests[2].py_num_accepted_draft_tokens == 4 + + +@pytest.mark.parametrize("budget", [1, 7, 8, 9, 64]) +def test_completion_budget_inside_a_full_block(budget): + """One token from the context step, then a fully accepted block of eight.""" + block = list(range(30, 38)) + (request,) = _run_block([budget], [block], [block[:7]]) + expected = ([9] + block)[:budget] + assert _generated(request) == expected + finished = request.state == LlmRequestState.GENERATION_COMPLETE + assert finished == (budget <= 9) + assert request.py_finish_reason == (FinishReason.LENGTH if finished else None) + + +# --------------------------------------------------------------------------------------------- +# The DFlash capture buffer +# --------------------------------------------------------------------------------------------- + +LAYERS = [4, 16, 28, 40, 46] +WIDTH = 5120 + + +def _tagged(layer: int, rows: int) -> torch.Tensor: + """Rows ``(layer, row)``-distinguishable in BF16: layer in column 0, row in column 1.""" + x = torch.zeros(rows, WIDTH, dtype=torch.bfloat16, device="cuda") + x[:, 0] = layer + x[:, 1] = torch.arange(rows, device="cuda", dtype=torch.bfloat16) + return x + + +def test_capture_slices_follow_the_layer_order(): + md = DFlashSpecMetadata( + max_draft_len=K, + max_total_draft_tokens=K, + spec_dec_mode=SpeculativeDecodingMode.DFLASH, + max_num_requests=4, + layers_to_capture=list(reversed(LAYERS)), + hidden_size=WIDTH, + max_num_tokens=64, + dtype=torch.bfloat16, + ) + for layer in range(52): + md.maybe_capture_hidden_states(layer, _tagged(layer, 24)) + captured = md.get_hidden_states(24) + assert captured.shape == (24, 5 * WIDTH) + for i, layer in enumerate(LAYERS): + assert torch.equal(captured[:, i * WIDTH : (i + 1) * WIDTH], _tagged(layer, 24)) + + +@pytest.mark.parametrize("offset", [-1, 1]) +def test_capture_at_a_neighbouring_layer_is_detected(offset): + """A target that taps layer l +/- 1 instead of l leaves a slice that is not layer l's.""" + md = DFlashSpecMetadata( + max_draft_len=K, + max_total_draft_tokens=K, + spec_dec_mode=SpeculativeDecodingMode.DFLASH, + max_num_requests=4, + layers_to_capture=LAYERS, + hidden_size=WIDTH, + max_num_tokens=64, + dtype=torch.bfloat16, + ) + md.captured_hidden_states.fill_(-1) + for layer in range(52): + md.maybe_capture_hidden_states(layer, _tagged(layer + offset, 16)) + captured = md.get_hidden_states(16) + wrong = [ + i + for i, layer in enumerate(LAYERS) + if not torch.equal(captured[:, i * WIDTH : (i + 1) * WIDTH], _tagged(layer, 16)) + ] + assert wrong == list(range(5)) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_serving.py b/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_serving.py new file mode 100644 index 000000000000..5f5d2958e6f0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_serving.py @@ -0,0 +1,359 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests of the Beam NVFP4 target's channel parser and family routing (generated inputs).""" + +import json +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2 import _router_index +from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget, routing_module +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.serving import ( # noqa: E501 + PARSER_KEY, + TAGS, + BeamHarmonyTextParser, +) +from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory + +pytestmark = pytest.mark.cpu_only + +#: (reply text after the generation prefix, expected content, expected reasoning_content) +REPLIES = { + "analysis_then_final": ( + "<|channel|>analysis<|message|>Add the two numbers.<|end|><|start|>assistant" + "<|channel|>final<|message|>2 + 3 = 5.<|return|>", + "2 + 3 = 5.", + "Add the two numbers.", + ), + "final_only": ("<|channel|>final<|message|>Paris.<|return|>", "Paris.", ""), + "no_headers": ("Plain text without channels.", "Plain text without channels.", ""), + "stopped_in_analysis": ("<|channel|>analysis<|message|>Still thinking", "", "Still thinking"), + "stopped_in_final": ( + "<|channel|>analysis<|message|>a<|end|><|start|>assistant<|channel|>final<|message|>Par", + "Par", + "a", + ), + "multiline_bodies": ( + "<|channel|>analysis<|message|>Line 1\n\nLine 2 x<|end|><|start|>assistant" + "<|channel|>final<|message|>\\(12 - 7 = 5\\)\n\nANSWER: 5<|return|>", + "\\(12 - 7 = 5\\)\n\nANSWER: 5", + "Line 1\n\nLine 2 x", + ), + "literal_angle_brackets": ( + "<|channel|>final<|message|>a <| b and <|x|><|return|>", + "a <| b and <|x|>", + "", + ), + "tool_channel_ignored": ( + "<|channel|>commentary to=functions.f <|constrain|>json<|message|>{}<|call|>", + "", + "", + ), + "terminated_by_call": ("<|channel|>final<|message|>done<|call|>ignored", "done", ""), +} + + +def _stream(text, pieces): + parser = BeamHarmonyTextParser() + content = reasoning = "" + for piece in pieces: + delta = parser.parse_delta(piece) + content += delta.content + reasoning += delta.reasoning_content + tail = parser.finish() + return content + tail.content, reasoning + tail.reasoning_content + + +@pytest.mark.parametrize("name", sorted(REPLIES)) +def test_whole_reply(name): + text, content, reasoning = REPLIES[name] + result = BeamHarmonyTextParser().parse(text) + assert (result.content, result.reasoning_content) == (content, reasoning) + assert not any(tag in result.content or tag in result.reasoning_content for tag in TAGS) + + +@pytest.mark.parametrize("name", sorted(REPLIES)) +def test_every_fragmentation_agrees_with_the_whole_reply(name): + text, content, reasoning = REPLIES[name] + assert _stream(text, list(text)) == (content, reasoning) + for cut in range(len(text) + 1): + assert _stream(text, [text[:cut], text[cut:]]) == (content, reasoning), cut + for size in (2, 3, 5, 7, 11): + pieces = [text[i : i + size] for i in range(0, len(text), size)] + assert _stream(text, pieces) == (content, reasoning), size + + +def test_parse_is_self_contained_and_repeatable(): + text, content, reasoning = REPLIES["analysis_then_final"] + parser = BeamHarmonyTextParser() + parser.parse_delta("<|channel|>analysis<|message|>left over") + for _ in range(2): + result = parser.parse(text) + assert (result.content, result.reasoning_content) == (content, reasoning) + + +def test_registered_under_its_key(): + parser = ReasoningParserFactory.create_reasoning_parser(PARSER_KEY) + assert isinstance(parser, BeamHarmonyTextParser) + assert ReasoningParserFactory.needs_raw_special_tokens(PARSER_KEY) + + +def _beam_config(**overrides): + fields = dict( + architectures=["BeamForCausalLM"], + num_hidden_layers=52, + hidden_size=5120, + num_experts=256, + quantization_config={"quant_algo": "NVFP4"}, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _mapping(**overrides): + fields = dict( + world_size=4, tp_size=4, pp_size=1, moe_tp_size=4, moe_ep_size=1, enable_attention_dp=False + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def test_route_claims_only_the_supported_deployment(): + route = routing_module("BeamForCausalLM").route + target = route(_beam_config(), _mapping(), (10, 3), None, None) + assert target.__name__ == "ModelingV2BeamFullModelNvfp4Sm103Tp4" + declined = [ + (_beam_config(), _mapping(), (10, 0), None, None), + (_beam_config(num_hidden_layers=6), _mapping(), (10, 3), None, None), + (_beam_config(quantization_config=None), _mapping(), (10, 3), None, None), + (_beam_config(), _mapping(tp_size=2, pp_size=2, moe_tp_size=2), (10, 3), None, None), + (_beam_config(), _mapping(moe_tp_size=1, moe_ep_size=4), (10, 3), None, None), + (_beam_config(), _mapping(enable_attention_dp=True), (10, 3), None, None), + (_beam_config(), _mapping(), (10, 3), object(), None), + (_beam_config(), _mapping(), (10, 3), None, "FP8"), + (_beam_config(), _mapping(), (10, 3), None, "NVFP4"), + ] + for args in declined: + with pytest.raises(NoTarget): + route(*args) + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +@pytest.mark.parametrize( + "checkpoint_kv,runtime_kv,supported", + [(None, "auto", True), (None, "fp8", False), (None, "nvfp4", False), ("FP8", "auto", False)], +) +def test_engine_resolver_honors_effective_kv_format( + monkeypatch, mode, checkpoint_kv, runtime_kv, supported +): + """The engine's complete routing call accepts BF16 KV and refuses quantized KV.""" + monkeypatch.setattr(_router_index, "_device_compute_capability", lambda: (10, 3)) + config = SimpleNamespace( + pretrained_config=_beam_config(), + mapping=_mapping(), + spec_config=None, + quant_config=SimpleNamespace(kv_cache_quant_algo=checkpoint_kv), + ) + args = SimpleNamespace(modeling_v2=mode, kv_cache_config=SimpleNamespace(dtype=runtime_kv)) + if supported: + target = _router_index.modeling_v2_resolve(config, args) + assert target.__name__ == "ModelingV2BeamFullModelNvfp4Sm103Tp4" + elif mode == "auto": + assert _router_index.modeling_v2_resolve(config, args) is None + else: + with pytest.raises(ValueError, match="KV cache"): + _router_index.modeling_v2_resolve(config, args) + + +@pytest.mark.parametrize("max_num_tokens", [0, -1, 8193, 16384]) +def test_launch_rejects_unsupported_token_budget_before_tokenizer(monkeypatch, max_num_tokens): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( + serving, + ) + + def unexpected_tokenizer(*args): + pytest.fail("invalid budget must be rejected before reading the checkpoint") + + monkeypatch.setattr(serving, "load_tokenizer", unexpected_tokenizer) + with pytest.raises(ValueError, match="max_num_tokens"): + serving.llm_args("unused", max_num_tokens=max_num_tokens) + + +@pytest.mark.parametrize("max_num_tokens", [1, 8192]) +def test_launch_accepts_token_budget_boundaries(monkeypatch, max_num_tokens): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( + serving, + ) + + tokenizer = object() + monkeypatch.setattr(serving, "load_tokenizer", lambda path: tokenizer) + args = serving.llm_args("unused", max_num_tokens=max_num_tokens) + assert args["max_num_tokens"] == max_num_tokens + assert args["tokenizer"] is tokenizer + + +def test_load_pretrained_config_preserves_beam_rope(tmp_path): + from tensorrt_llm._torch.models.modeling_beam import BeamConfig + from tensorrt_llm._torch.pyexecutor.config_utils import load_pretrained_config + + rope = { + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + } + config = { + "model_type": "beam", + "num_hidden_layers": 6, + "torch_dtype": "bfloat16", + "rope_parameters": rope, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + loaded = load_pretrained_config(str(tmp_path)) + assert isinstance(loaded, BeamConfig) + assert loaded.name_or_path == str(tmp_path) + assert loaded.torch_dtype == torch.bfloat16 + assert loaded.rope_parameters == rope + + +# --------------------------------------------------------------------------------------------- +# DFlash 2 speculative decoding: routing, the launcher's config and the drafter contract +# --------------------------------------------------------------------------------------------- + + +def _dflash(**overrides): + from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig + + fields = dict(max_draft_len=7, attention_backend="TRTLLM", speculative_model="/drafter") + fields.update(overrides) + return DFlashDecodingConfig(**fields) + + +def _mtp(): + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + + return MTPDecodingConfig(max_draft_len=3) + + +def test_route_claims_the_dflash_deployment(): + route = routing_module("BeamForCausalLM").route + target = route(_beam_config(), _mapping(), (10, 3), _dflash(), None) + assert target.__name__ == "ModelingV2BeamFullModelNvfp4Sm103Tp4" + args = SimpleNamespace(disable_overlap_scheduler=True) + assert route(_beam_config(), _mapping(), (10, 3), _dflash(), None, llm_args=args) is target + + +@pytest.mark.parametrize( + "spec,args,reason", + [ + (lambda: _dflash(attention_backend="VANILLA"), None, "only the TRTLLM backend reads"), + (lambda: _dflash(attention_backend="FA4"), None, "only the TRTLLM backend reads"), + (lambda: _dflash(max_draft_len=15), None, "max_draft_len 15, not 7"), + (lambda: _dflash(max_draft_len=3), None, "max_draft_len 3, not 7"), + (lambda: _dflash(draft_len_schedule={1: 7, 4: 3}), None, "draft_len_schedule"), + (lambda: _dflash(use_rejection_sampling=True), None, "rejection sampling"), + (lambda: _dflash(speculative_model=None), None, "drafter checkpoint"), + ( + lambda: _dflash(), + SimpleNamespace(disable_overlap_scheduler=False), + "overlap scheduler", + ), + (lambda: _mtp(), None, "is not DFlash"), + ], + ids=["vanilla", "fa4", "k15", "k3", "schedule", "rejection", "no-drafter", "overlap", "mtp"], +) +def test_route_refuses_other_speculation_clearly(spec, args, reason): + route = routing_module("BeamForCausalLM").route + with pytest.raises(NoTarget, match=reason): + route(_beam_config(), _mapping(), (10, 3), spec(), None, llm_args=args) + + +def test_launcher_dflash_config(tmp_path): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( # noqa: E501 + serving, + ) + + with pytest.raises(ValueError, match="has no config.json"): + serving.dflash_config(str(tmp_path / "missing")) + (tmp_path / "config.json").write_text("{}") + config = serving.dflash_config(str(tmp_path)) + assert (config.max_draft_len, config.attention_backend) == (7, "TRTLLM") + assert config.speculative_model == str(tmp_path) + assert not config.use_rejection_sampling + + +def test_launcher_args_add_speculation_only_with_a_drafter(tmp_path, monkeypatch): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( # noqa: E501 + serving, + ) + + monkeypatch.setattr(serving, "load_tokenizer", lambda checkpoint: "tokenizer") + (tmp_path / "config.json").write_text("{}") + plain = serving.llm_args("/target", max_seq_len=8192, max_batch_size=4) + assert "speculative_config" not in plain + spec = serving.llm_args("/target", max_seq_len=8192, draft_checkpoint=str(tmp_path)) + assert spec["speculative_config"].max_draft_len == 7 + assert {k: spec[k] for k in serving.DEPLOYMENT} == serving.DEPLOYMENT + assert spec["cuda_graph_config"] is None and spec["disable_overlap_scheduler"] + + +def _drafter_config(**dflash_overrides): + dflash = dict( + target_layer_ids=[4, 16, 28, 40, 46], + conv_kernel_size=2, + conv_group_size=16, + selector_rank=256, + selector_top_k=16, + final_logit_softcapping=30, + ) + dflash.update(dflash_overrides) + pretrained = SimpleNamespace( + hidden_size=5120, vocab_size=151936, draft_vocab_size=151936, dflash_config=dflash + ) + return SimpleNamespace(pretrained_config=pretrained) + + +_TARGET_CONFIG = {"hidden_size": 5120, "vocab_size": 151936} + + +def test_drafter_contract_accepts_the_dflash2_drafter(): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.modeling import ( # noqa: E501 + check_drafter_config, + ) + + check_drafter_config(_drafter_config(), _TARGET_CONFIG, [4, 16, 28, 40, 46]) + + +@pytest.mark.parametrize( + "config,layers,reason", + [ + (_drafter_config(), [3, 15, 27, 39, 45], "target_layer_ids"), + (_drafter_config(target_layer_ids=[4, 16, 28, 40, 52]), [4, 16, 28, 40, 52], "0..51"), + (_drafter_config(target_layer_ids=[16, 4]), [16, 4], "ascending"), + (_drafter_config(final_logit_softcapping=None), [4, 16, 28, 40, 46], "softcapping"), + (_drafter_config(selector_rank=0), [4, 16, 28, 40, 46], "selector_rank"), + ], + ids=["shifted", "out-of-range", "unordered", "no-softcap", "no-selector"], +) +def test_drafter_contract_refuses_a_mismatch(config, layers, reason): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.modeling import ( # noqa: E501 + check_drafter_config, + ) + + with pytest.raises(ValueError, match=reason): + check_drafter_config(config, _TARGET_CONFIG, layers) + + +def test_drafter_contract_refuses_another_vocabulary_or_width(): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.modeling import ( # noqa: E501 + check_drafter_config, + ) + + config = _drafter_config() + config.pretrained_config.draft_vocab_size = 32000 + with pytest.raises(ValueError, match="draft_vocab_size 32000"): + check_drafter_config(config, _TARGET_CONFIG, [4, 16, 28, 40, 46]) + config = _drafter_config() + config.pretrained_config.hidden_size = 4096 + with pytest.raises(ValueError, match="hidden_size 4096"): + check_drafter_config(config, _TARGET_CONFIG, [4, 16, 28, 40, 46]) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1/_moe_reference.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1/_moe_reference.py new file mode 100644 index 000000000000..aa0029b67cfa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1/_moe_reference.py @@ -0,0 +1,287 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plain-Torch reference for the MoE feed-forward of Beam decoder block 3. + +An independent restatement of the source MoE for generated-input comparison: + +* Routing, every step fp32 with TF32 off: logits ``x.float() @ gate.weight^T``, sigmoid + scores, and selection scores = scores + ``gate.e_score_correction_bias``, which only + choose. Each of the 8 groups of 32 experts is scored by the sum of its two largest + selection scores; the 4 best groups are kept, every expert outside them is excluded, + and the 8 best remaining selection scores are chosen in descending order. The mixing + weights are the chosen experts' unbiased scores over their sum plus 1e-20. +* Each selected expert, in ascending id order, runs its bias-free bf16 SwiGLU on its + tokens' rows in ascending token order: bf16 gate and up projections, fp32 SiLU and + product, one bf16 rounding, bf16 down projection. +* Each expert output times its fp32 weight is rounded to bf16, the routes are summed per + token in fp32, and the sum is scaled by the routed scaling factor 2.5 once and rounded + to bf16. + +Only torch is imported. Group exclusion is written as group membership; the two +selections stay ``topk(..., sorted=True)``, the source's own calls, because the order of +tied scores is whatever that call returns. + +State and inputs are generated on the GPU. Every tensor draws from a generator seeded +with ``crc32(f"{seed}|{name}")``, so its values do not depend on which other tensors were +generated or in which order. A device generator's draws depend on the device and the +torch build: the bytes are reproducible there, not across them. +""" + +from __future__ import annotations + +import contextlib +import math +import zlib +from collections import OrderedDict +from typing import Callable, Iterator + +import torch +import torch.nn.functional as F + +HIDDEN = 5120 +INTERMEDIATE = 2560 +NUM_EXPERTS = 256 +N_GROUP = 8 +PER_GROUP = NUM_EXPERTS // N_GROUP +TOPK_GROUP = 4 +TOP_K = 8 +NORM_EPS = 1e-20 +ROUTED_SCALING_FACTOR = 2.5 +MAX_TOKENS = 129 +SEEDS = (17, 41) +LENGTHS = (1, 17, 129) + +#: Each expert's three tensors by name relative to ``experts..``, [out, in]. +EXPERT_SHAPES = OrderedDict( + [ + ("gate_proj.weight", (INTERMEDIATE, HIDDEN)), + ("up_proj.weight", (INTERMEDIATE, HIDDEN)), + ("down_proj.weight", (HIDDEN, INTERMEDIATE)), + ] +) + +STATE_RECIPE = { + "seeding": "one torch.Generator on the device per call, re-seeded for every tensor with " + "zlib.crc32(f'{seed}|{name}')", + "gate.weight": "fp32 N(0, 1) / sqrt(5120)", + "gate.e_score_correction_bias": "fp32 0.05 * N(0, 1)", + "experts..gate_proj.weight, experts..up_proj.weight": "N(0, 1) / sqrt(5120), drawn " + "in fp32 and rounded once to bf16", + "experts..down_proj.weight": "N(0, 1) / sqrt(2560), drawn in fp32 and rounded once to bf16", + "input": "N(0, 1) of shape [129, 5120] under the name 'input', drawn in fp32 and rounded " + "once to bf16; a T-token case uses the first T rows", +} + +#: ``record(name, value)``: receives the named intermediates of a call. +Record = Callable[[str, object], None] + + +# ----------------------------------------------------------------------------- numerics + + +@contextlib.contextmanager +def reference_math() -> Iterator[None]: + """IEEE fp32 GEMMs (TF32 off), bf16 GEMMs without reduced-precision reductions and + autocast off for the duration; the caller's settings are restored afterwards.""" + matmul = torch.backends.cuda.matmul + saved_tf32 = matmul.allow_tf32 + saved_precision = getattr(matmul, "fp32_precision", None) + saved_bf16 = matmul.allow_bf16_reduced_precision_reduction + matmul.allow_tf32 = False + if saved_precision is not None: + # Both switches: with only the legacy one moved, fp32_precision stays "tf32". + matmul.fp32_precision = "ieee" + matmul.allow_bf16_reduced_precision_reduction = False + try: + with torch.autocast("cuda", enabled=False): + yield + finally: + # The legacy switch first: its setter also moves the newer one restored after it. + matmul.allow_tf32 = saved_tf32 + if saved_precision is not None: + matmul.fp32_precision = saved_precision + matmul.allow_bf16_reduced_precision_reduction = saved_bf16 + + +# ----------------------------------------------------------------------------- state and inputs + + +def tensor_seed(seed: int, name: str) -> int: + """The generator seed of the tensor called ``name`` in case ``seed``.""" + return zlib.crc32(f"{seed}|{name}".encode()) + + +def param_shapes() -> "OrderedDict[str, tuple[tuple[int, ...], torch.dtype]]": + """The MoE's 770 tensors by checkpoint name relative to ``mlp.``, in checkpoint order: + the router's fp32 weight and correction bias, then each expert's three bf16 matrices.""" + shapes = OrderedDict( + [ + ("gate.weight", ((NUM_EXPERTS, HIDDEN), torch.float32)), + ("gate.e_score_correction_bias", ((NUM_EXPERTS,), torch.float32)), + ] + ) + for expert_id in range(NUM_EXPERTS): + for name, shape in EXPERT_SHAPES.items(): + shapes[f"experts.{expert_id}.{name}"] = (shape, torch.bfloat16) + return shapes + + +def _draw( + gen: torch.Generator, seed: int, name: str, shape: tuple[int, ...], device +) -> torch.Tensor: + """Standard normal fp32 values of ``shape`` from ``gen`` re-seeded for ``name``.""" + gen.manual_seed(tensor_seed(seed, name)) + return torch.randn(shape, generator=gen, dtype=torch.float32, device=device) + + +def generate_state( + seed: int, device: str | torch.device = "cuda" +) -> "OrderedDict[str, torch.Tensor]": + """The seed's 770 tensors on ``device``, in checkpoint order (see `STATE_RECIPE`). + + Every matrix is scaled by 1 / sqrt(its input width), the correction bias by 0.05; the + expert matrices are rounded once to bf16, the router stays fp32. + """ + gen = torch.Generator(device=device) + state: "OrderedDict[str, torch.Tensor]" = OrderedDict() + for name, (shape, dtype) in param_shapes().items(): + scale = 0.05 if len(shape) == 1 else 1.0 / math.sqrt(shape[1]) + state[name] = (_draw(gen, seed, name, shape, device) * scale).to(dtype).contiguous() + return state + + +def generate_input( + seed: int, tokens: int = MAX_TOKENS, device: str | torch.device = "cuda" +) -> torch.Tensor: + """``[tokens, 5120]`` bf16: the first ``tokens`` rows of the seed's 129-row input.""" + if not 1 <= tokens <= MAX_TOKENS: + raise ValueError(f"tokens must be in 1..{MAX_TOKENS}, got {tokens}") + gen = torch.Generator(device=device) + rows = _draw(gen, seed, "input", (MAX_TOKENS, HIDDEN), device).to(torch.bfloat16) + return rows[:tokens] + + +def state_problems(state: "dict[str, torch.Tensor]", seed: int) -> list[str]: + """Exact accounting plus the properties the comparisons rely on: finite values, a bias + that varies across experts, and no two expert matrices alike.""" + problems = [] + shapes = param_shapes() + if list(state) != list(shapes): + problems.append("names or order differ from param_shapes()") + if len({tensor_seed(seed, name) for name in shapes}) != len(shapes): + problems.append("two tensors share a generator seed") + for name, (shape, dtype) in shapes.items(): + tensor = state.get(name) + if tensor is None: + continue + if tuple(tensor.shape) != shape or tensor.dtype != dtype: + found = f"{tuple(tensor.shape)} {tensor.dtype}" + problems.append(f"{name}: {found}, expected {shape} {dtype}") + elif not bool(torch.isfinite(tensor).all()): + problems.append(f"{name}: nonfinite") + if problems: + return problems + bias = state["gate.e_score_correction_bias"] + if bool((bias == bias[0]).all()): + problems.append("the correction bias does not vary across experts") + for width, names in ( + (HIDDEN, ("gate_proj.weight", "up_proj.weight")), + (INTERMEDIATE, ("down_proj.weight",)), + ): + first_rows = torch.stack( + [state[f"experts.{e}.{n}"][0].float() for e in range(NUM_EXPERTS) for n in names] + ) + if torch.unique(first_rows, dim=0).shape[0] != first_rows.shape[0]: + problems.append(f"two expert matrices of input width {width} share their first row") + return problems + + +# ----------------------------------------------------------------------------- the MoE + + +def route( + x: torch.Tensor, weight: torch.Tensor, correction_bias: torch.Tensor +) -> dict[str, torch.Tensor]: + """The source routing of ``x [T, 5120]`` bf16, every step fp32. + + Returns ``logits``, ``scores`` and ``selection`` ``[T, 256]``, ``group_scores`` + ``[T, 8]``, ``groups`` ``[T, 4]`` int64 (descending), ``masked`` ``[T, 256]`` (-inf + outside the kept groups), ``ids`` ``[T, 8]`` int64 (descending selection score) and + ``weights`` ``[T, 8]`` fp32. + """ + with reference_math(): + logits = F.linear(x.float(), weight.float()) + scores = torch.sigmoid(logits) + selection = scores + correction_bias.float() + tokens = selection.shape[0] + group_scores = selection.view(tokens, N_GROUP, PER_GROUP).topk(2, dim=-1).values.sum(dim=-1) + groups = group_scores.topk(TOPK_GROUP, dim=-1, sorted=True).indices + expert_group = torch.arange(NUM_EXPERTS, device=x.device) // PER_GROUP + kept = (expert_group.view(1, NUM_EXPERTS, 1) == groups.unsqueeze(1)).any(dim=-1) + masked = torch.where(kept, selection, torch.full_like(selection, float("-inf"))) + ids = masked.topk(TOP_K, dim=-1, sorted=True).indices + chosen = torch.gather(scores, -1, ids) + weights = chosen / (chosen.sum(dim=-1, keepdim=True) + NORM_EPS) + return dict( + logits=logits, + scores=scores, + selection=selection, + group_scores=group_scores, + groups=groups, + masked=masked, + ids=ids, + weights=weights, + ) + + +def expert( + rows: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> dict[str, torch.Tensor]: + """One expert on ``rows [n, 5120]`` bf16, every step: ``gate``, ``up`` (bf16), + ``silu``, ``product`` (fp32), ``act`` (bf16, the one rounding) and ``y`` ``[n, 5120]`` + bf16.""" + with reference_math(): + gate = F.linear(rows, gate_weight) + up = F.linear(rows, up_weight) + silu = F.silu(gate.float()) + product = silu * up.float() + act = product.to(torch.bfloat16) + y = F.linear(act, down_weight) + return dict(gate=gate, up=up, silu=silu, product=product, act=act, y=y) + + +def moe( + x: torch.Tensor, state: "dict[str, torch.Tensor]", record: Record | None = None +) -> torch.Tensor: + """The source MoE of ``x [T, 5120]`` bf16; returns ``[T, 5120]`` bf16. + + ``state`` holds the 770 tensors by name relative to ``mlp.``. ``record``, when given, + receives ``routing`` (the dict of `route`), ``expert_calls`` (``[(expert_id, tokens)]`` + in call order), ``accumulator`` (the fp32 sum before scaling) and ``output``. + """ + assert x.dtype == torch.bfloat16 and x.dim() == 2 and x.shape[1] == HIDDEN, ( + x.dtype, + tuple(x.shape), + ) + routing = route(x, state["gate.weight"], state["gate.e_score_correction_bias"]) + ids, weights = routing["ids"], routing["weights"] + acc = torch.zeros(x.shape[0], HIDDEN, dtype=torch.float32, device=x.device) + calls = [] + for expert_id in sorted(set(ids.reshape(-1).tolist())): + # Row-major nonzero: the expert's tokens in ascending order, each once. + tokens, slots = (ids == expert_id).nonzero(as_tuple=True) + matrices = [state[f"experts.{expert_id}.{name}"] for name in EXPERT_SHAPES] + y = expert(x[tokens], *matrices)["y"] + routed = (y.float() * weights[tokens, slots].unsqueeze(-1)).to(torch.bfloat16) + acc[tokens] += routed.float() + calls.append((expert_id, tokens)) + output = (acc * ROUTED_SCALING_FACTOR).to(torch.bfloat16) + if record is not None: + record("routing", routing) + record("expert_calls", calls) + record("accumulator", acc) + record("output", output) + return output diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1/test_moe.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1/test_moe.py new file mode 100644 index 000000000000..117bbe55508b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1/test_moe.py @@ -0,0 +1,1179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the Beam MoE component of decoder block 3: bf16 / sm_103 / tp1. + +Configuration under test, fixed below and nowhere else: block 3's MoE at hidden 5120 -- +256 experts of intermediate 2560 in 8 groups of 32, 4 groups and 8 experts per token, +routed scaling factor 2.5; an fp32 router and bf16 experts with the source model's fp32 +steps; one sm_103 GPU; eager, one call per case on already-normalized bf16 hidden +states; generated weights; seeds {17, 41} x lengths {1, 17, 129}, and designed fixtures. + +What is compared. ``_moe_reference.py`` restates the source MoE in plain torch and +generates its state and input deterministically on the GPU. For every case the forward's +own routing must select exactly the reference's groups and experts, in order, and its +router steps and its output are held to the rules in `RULES`. The forward's op and +expert calls, and fixtures whose routing is fixed by design, show that each token +reaches its own experts and that expert ``e`` is the one loaded under ``experts.``. + +Invocation (fp32 steps need the TF32 overrides removed before torch is imported):: + + env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE \\ + python -m pytest tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1 + +Optional environment: + +* ``BEAM_MOE_EVIDENCE``: a directory that receives one JSON record per compared case: + every comparison with its rule, the op and expert calls where they were recorded, and + the device and software it ran on; and ``fp32-session.json``, the fp32 GEMM checks made + before and after the module's tests. + +Dependencies: a TensorRT-LLM build for sm_103 and about 60 GiB of free device memory: +one seed's generated state and the loaded MoE hold 18.75 GiB of expert weights each, and +one case builds a second, unloaded MoE. No model files. +""" + +from __future__ import annotations + +import contextlib +import gc +import hashlib +import json +import os +import platform +import subprocess +import sys +import types +from collections import OrderedDict +from pathlib import Path + +import pytest +import torch + +import tensorrt_llm +import tensorrt_llm._torch.models.modeling_beam as builtin_beam +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_bf16__sm_103__tp1 import ( + moe as moe_module, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_bf16__sm_103__tp1.moe import ( + BeamMoeBf16Sm103Tp1, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_dispatch_combine_bf16__sm_103__tp1 import ( # noqa: E501 + dispatch_combine, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_dispatch_combine_bf16__sm_103__tp1.ops.moe.beam_dispatch_combine import ( # noqa: E501 + beam_dispatch_combine_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_expert_bf16__sm_103__tp1.expert import ( # noqa: E501 + BeamMoeExpertBf16Sm103Tp1, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_expert_bf16__sm_103__tp1.ops.activation.swiglu_fp32 import ( # noqa: E501 + swiglu_fp32, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_router_bf16__sm_103__tp1.ops.moe.beam_grouped_topk_router import ( # noqa: E501 + beam_grouped_topk_router, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_router_bf16__sm_103__tp1.router import ( # noqa: E501 + BeamMoeRouterBf16Sm103Tp1, +) + +__extra_import_path__ = [".", "../../.."] +import _moe_reference as reference # noqa: E402 +from _validating import validating # noqa: E402 + +assert torch.cuda.is_available(), "the Beam MoE tests require a CUDA device" + +pytestmark = pytest.mark.high_cuda_memory + +# -------------------------------------------------- the fixed configuration + +#: Block 3's MoE as the checkpoint configures it. Only these fields are read. +MOE_CONFIG = types.MappingProxyType( + dict( + hidden_size=5120, + moe_intermediate_size=2560, + num_experts=256, + num_experts_per_tok=8, + n_group=8, + topk_group=4, + norm_topk_prob=True, + routed_scaling_factor=2.5, + hidden_act="silu", + mlp_layer_types=("dense", "dense", "dense", "sparse"), + ) +) +LAYER_IDX = 3 +DTYPE = torch.bfloat16 +SM = (10, 3) +SEEDS = (17, 41) +LENGTHS = (1, 17, 129) +CASES = tuple((seed, length) for seed in SEEDS for length in LENGTHS) +CASE_IDS = [f"seed{seed}-T{length}" for seed, length in CASES] +#: Eager, one GPU, one call per case on already-normalized hidden states. +EXECUTION = types.MappingProxyType(dict(eager=True, cuda_graph=False, devices=1)) + +# -------------------------------------------------- comparison rules + +_RTOL = 2.0**-7 + +#: (atol, rtol, dtype) per compared tensor. An element passes when +#: |actual - reference| <= atol + rtol * |reference|, both widened to fp64; every element +#: of every compared tensor must pass, and no NaN or Inf is accepted. rtol 2^-7 is one bf16 +#: unit at any magnitude. Routing is never excused by a tolerance: the groups and the +#: expert ids must equal the reference's exactly, in order. +RULES = types.MappingProxyType( + { + "moe_output": (0.03, _RTOL, torch.bfloat16), + "router_logits": (2e-5, 0.0, torch.float32), + "router_scores": (3e-6, 0.0, torch.float32), + "router_selection_scores": (3e-6, 0.0, torch.float32), + "router_weights": (8e-7, 0.0, torch.float32), + } +) +#: Per token, |sum(mixing weights) - 1| may not exceed this and no weight may reach +#: `MAX_WEIGHT`: the routed scaling factor 2.5 folded into the weights breaks both. +WEIGHT_SUM_ATOL = 1e-6 +MAX_WEIGHT = 1.5 + +#: Every catalog entry the MoE's forward can call. +OPS = (beam_grouped_topk_router, cublas_mm, swiglu_fp32, beam_dispatch_combine_op) +#: What each selected expert runs, in order, inside BeamDispatchCombine. +EXPERT_CALLS = ("expert", "CublasMm", "CublasMm", "SwigluFp32", "CublasMm") +#: Keyword arguments of the entries' calls that identify the path. +_RECORDED_KWARGS = ("n_group", "topk_group", "top_k", "eps", "routed_scaling_factor") + +# -------------------------------------------------- designed routing + +#: Two sets of eight experts, the first and last expert of four groups each: A holds groups +#: 0, 1, 4 and 7, both ends of the id range included; B holds groups 2, 3, 5 and 6. +SET_A = (0, 31, 32, 63, 128, 159, 224, 255) +SET_B = (64, 95, 96, 127, 160, 191, 192, 223) +#: Per designed token, the rank of each member of its set, in the set's order. Rank r is the +#: logit -1 + r / 2; every other expert's logit is -2. +RANKS = ( + (0, 1, 2, 3, 4, 5, 6, 7), + (7, 6, 5, 4, 3, 2, 1, 0), + (3, 7, 1, 5, 0, 4, 2, 6), + (6, 2, 4, 0, 5, 1, 7, 3), +) +#: The hidden column that carries each designed token's logits, per set. +COLUMNS_A = tuple(5 + t * (5120 // 4) for t in range(len(RANKS))) +COLUMNS_B = tuple(5 + 5120 // 8 + t * (5120 // 4) for t in range(len(RANKS))) +SPIKE = 16.0 +OTHER_LOGIT = -2.0 +FIXTURE_SEED = 17 + + +def designed_logit(rank: int) -> float: + return -1.0 + 0.5 * rank + + +def moe_config(**overrides) -> types.SimpleNamespace: + fields = dict(MOE_CONFIG) + fields["mlp_layer_types"] = list(fields["mlp_layer_types"]) + fields.update(overrides) + return types.SimpleNamespace(**fields) + + +# -------------------------------------------------- precision and environment + + +def _fp32_gemm_probe() -> float: + """Worst relative error of an fp32 GEMM against fp64; TF32 substitution shows as ~1e-3.""" + gen = torch.Generator().manual_seed(1234) + x = torch.randn(64, 5120, generator=gen).to(torch.bfloat16).float().cuda() + w = (torch.randn(512, 5120, generator=gen) * 5120**-0.5).float().cuda() + y = torch.nn.functional.linear(x, w) + ref = x.double() @ w.double().t() + return float((y.double() - ref).abs().max()) / float(ref.pow(2).mean().sqrt()) + + +def _math_state() -> dict: + return { + "float32_matmul_precision": torch.get_float32_matmul_precision(), + "cuda_matmul_allow_tf32": bool(torch.backends.cuda.matmul.allow_tf32), + "cudnn_allow_tf32": bool(torch.backends.cudnn.allow_tf32), + } + + +@pytest.fixture(scope="module", autouse=True) +def ieee_fp32(): + """The source's fp32 steps are IEEE fp32: refuse a process that substitutes TF32, both + before the module's tests and after them.""" + for var in ("TORCH_ALLOW_TF32_CUBLAS_OVERRIDE", "NVIDIA_TF32_OVERRIDE"): + assert var not in os.environ, ( + f"{var} is set and substitutes TF32 for fp32 GEMMs; run with `env -u {var}`" + ) + saved = ( + torch.get_float32_matmul_precision(), + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + ) + torch.set_float32_matmul_precision("highest") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + before = _fp32_gemm_probe() + assert before <= 1e-5, f"fp32 GEMM is not IEEE fp32 (relative error {before:.3g})" + state = {"math": _math_state(), "fp32_gemm_probe_before": before} + yield state + after = _fp32_gemm_probe() + state.update(math_after=_math_state(), fp32_gemm_probe_after=after) + _write_evidence("fp32-session", {"fixture": "ieee_fp32", "math": state}) + precision, allow_matmul, allow_cudnn = saved + torch.set_float32_matmul_precision(precision) + torch.backends.cuda.matmul.allow_tf32 = allow_matmul + torch.backends.cudnn.allow_tf32 = allow_cudnn + assert after <= 1e-5, f"fp32 GEMM stopped being IEEE fp32 (relative error {after:.3g})" + + +@pytest.fixture(scope="module", autouse=True) +def device_memory(): + """Hand the module's cached state and MoE back to the device when its tests are done.""" + yield + release_cache() + + +def _git_revision() -> str | None: + root = Path(tensorrt_llm.__file__).resolve().parents[1] + try: + out = subprocess.run( + ["git", "-C", str(root), "rev-parse", "HEAD"], + capture_output=True, + text=True, + timeout=30, + ) + except (OSError, subprocess.SubprocessError): + return None + return out.stdout.strip() or None + + +def _environment(ieee: dict) -> dict: + device = torch.cuda.current_device() + return { + "device_name": torch.cuda.get_device_name(device), + "device_capability": list(torch.cuda.get_device_capability(device)), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "tensorrt_llm": tensorrt_llm.__version__, + "tensorrt_llm_file": tensorrt_llm.__file__, + "source_revision": _git_revision(), + "python": sys.version.split()[0], + "machine": platform.machine(), + "argv": sys.argv, + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "slurm_step_id": os.environ.get("SLURM_STEP_ID"), + "execution": dict(EXECUTION), + "math": ieee, + } + + +def _write_evidence(name: str, record: dict) -> None: + directory = os.environ.get("BEAM_MOE_EVIDENCE") + if not directory: + return + path = Path(directory) + path.mkdir(parents=True, exist_ok=True) + (path / f"{name}.json").write_text(json.dumps(record, indent=1, default=str) + "\n") + + +def tensor_sha256(t: torch.Tensor) -> str: + flat = t.detach().contiguous().cpu().reshape(-1) + h = hashlib.sha256() + h.update(f"{t.dtype}|{tuple(t.shape)}|".encode()) + h.update(flat.view(torch.uint8).numpy().tobytes()) + return h.hexdigest() + + +# -------------------------------------------------- comparison + + +def compare(name: str, actual: torch.Tensor, expected: torch.Tensor, rule: str | None) -> dict: + """Score one tensor against one rule; `rule=None` measures without judging.""" + result = { + "tensor": name, + "rule": rule, + "shape": list(expected.shape), + "dtype": str(actual.dtype), + } + problems = [] + dtype = expected.dtype if rule is None else RULES[rule][2] + if tuple(actual.shape) != tuple(expected.shape): + problems.append(f"shape {tuple(actual.shape)} != {tuple(expected.shape)}") + if actual.dtype != dtype or expected.dtype != dtype: + problems.append(f"dtypes {actual.dtype}/{expected.dtype}, need {dtype}") + if problems: + result.update(passed=False, problems=problems) + return result + a, r = actual.double(), expected.to(actual.device).double() + if not bool(torch.isfinite(a).all()) or not bool(torch.isfinite(r).all()): + result.update(passed=False, problems=["non-finite values"]) + return result + diff = (a - r).abs() + numel = diff.numel() + result.update(numel=numel, unequal=int((a != r).sum()), max_abs=0.0) + if numel: + worst = int(diff.argmax()) + result.update( + max_abs=float(diff.max()), + mean_abs=float(diff.mean()), + worst_index=[int(i) for i in torch.unravel_index(torch.tensor(worst), diff.shape)], + worst_actual=float(a.reshape(-1)[worst]), + worst_expected=float(r.reshape(-1)[worst]), + ) + if rule is None: + result["passed"] = None + return result + atol, rtol, _ = RULES[rule] + violations = int((diff > atol + rtol * r.abs()).sum()) + result.update(atol=atol, rtol=rtol, violations=violations) + if violations: + problems.append( + f"{violations} of {numel} elements outside atol {atol} + rtol {rtol} * |ref| " + f"(max abs {result['max_abs']:.3g})" + ) + result.update(passed=not problems, problems=problems) + return result + + +def compare_exact(name: str, actual: torch.Tensor, expected: torch.Tensor) -> dict: + """Integer selections: equal element for element and in order, never sorted first.""" + same_kind = tuple(actual.shape) == tuple(expected.shape) and actual.dtype == expected.dtype + if same_kind: + rows = (actual == expected).reshape(expected.shape[0], -1).all(dim=-1) + mismatched = [int(i) for i in torch.nonzero(~rows).flatten().cpu()] + else: + mismatched = list(range(expected.shape[0])) + problems = [] + if not same_kind: + problems.append( + f"{tuple(actual.shape)} {actual.dtype} != {tuple(expected.shape)} {expected.dtype}" + ) + elif mismatched: + problems.append(f"rows {mismatched[:8]} differ") + return { + "tensor": name, + "rule": "exact, in order", + "shape": list(expected.shape), + "mismatched_rows": mismatched, + "passed": not problems, + "problems": problems, + } + + +def compare_routes(route: dict, ref: dict) -> list[dict]: + """The router op's mapping of one call against the reference route of the same input.""" + results = [ + compare_exact("groups", route["groups"], ref["groups"]), + compare_exact("ids", route["ids"], ref["ids"]), + compare("router_logits", route["logits"], ref["logits"], "router_logits"), + compare("router_scores", route["scores"], ref["scores"], "router_scores"), + compare( + "router_selection_scores", + route["selection_scores"], + ref["selection"], + "router_selection_scores", + ), + compare("router_group_scores", route["group_scores"], ref["group_scores"], None), + compare("router_weights", route["weights"], ref["weights"], "router_weights"), + ] + same = torch.equal(torch.isinf(route["masked_selection_scores"]), torch.isinf(ref["masked"])) + results.append( + { + "tensor": "excluded_experts", + "rule": "exactly the reference's -inf exclusions", + "passed": same, + "problems": [] if same else ["the excluded experts differ"], + } + ) + return results + + +def weight_invariants(route: dict) -> dict: + """Mixing weights sum to one per token and carry no routed scaling factor.""" + weights = route["weights"] + deviation = float((weights.double().sum(dim=-1) - 1.0).abs().max()) + largest = float(weights.max()) + problems = [] + if deviation > WEIGHT_SUM_ATOL: + problems.append(f"a token's weights sum to 1 +- {deviation:.3g}") + if largest >= MAX_WEIGHT: + problems.append(f"a weight of {largest:.4g}") + return { + "tensor": "mixing_weights", + "rule": f"|sum - 1| <= {WEIGHT_SUM_ATOL:g} per token, every weight < {MAX_WEIGHT}", + "max_abs_sum_minus_one": deviation, + "max_weight": largest, + "passed": not problems, + "problems": problems, + } + + +def output_contract(out, tokens: int) -> dict: + problems = [] + if not isinstance(out, torch.Tensor): + problems.append(f"returned {type(out).__name__}") + elif tuple(out.shape) != (tokens, 5120) or out.dtype != DTYPE or not out.is_cuda: + problems.append(f"returned {tuple(out.shape)} {out.dtype} on {out.device}") + elif torch.cuda.get_device_capability(out.device) != SM: + problems.append(f"on {torch.cuda.get_device_name(out.device)}") + elif not bool(torch.isfinite(out).all()): + problems.append("non-finite values") + return { + "tensor": "output_contract", + "rule": f"[{tokens}, 5120] bf16 on the sm_103 device, finite", + "passed": not problems, + "problems": problems, + } + + +def tensor_versions(module: torch.nn.Module) -> dict[str, int]: + """The in-place write counter of every parameter and buffer, by name.""" + versions = {name: p._version for name, p in module.named_parameters()} + versions.update({name: b._version for name, b in module.named_buffers()}) + return versions + + +def unchanged(x: torch.Tensor, x_before: torch.Tensor, before: dict, after: dict) -> dict: + written = sorted(name for name, version in before.items() if after.get(name) != version) + problems = [] + if not bitwise_equal(x, x_before): + problems.append("the input was modified") + if written: + problems.append(f"{len(written)} parameters or buffers were written: {written[:4]}") + return { + "tensor": "input_and_parameters", + "rule": "input bits and parameter write counters unchanged by the call", + "passed": not problems, + "problems": problems, + } + + +def bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + """Same dtype, shape and bits (so -0.0 differs from +0.0).""" + if a.dtype != b.dtype or a.shape != b.shape: + return False + view = {torch.bfloat16: torch.int16, torch.float32: torch.int32}.get(a.dtype) + if view is None: + return bool(torch.equal(a, b)) + return bool(torch.equal(a.contiguous().view(view), b.contiguous().view(view))) + + +def assert_all_pass(results: list[dict]) -> None: + failed = [r for r in results if r["passed"] is False] + assert not failed, "\n".join( + f"{r['tensor']} [{r['rule']}]: {'; '.join(r['problems'])}" for r in failed + ) + + +# -------------------------------------------------- state, inputs and the loaded MoE + +_STATE_CACHE: dict[int, dict] = {} + + +def release_cache() -> None: + """Drop the cached state and MoE and hand their memory back to the device.""" + _STATE_CACHE.clear() + gc.collect() + torch.cuda.empty_cache() + + +def generated(seed: int) -> dict: + """The seed's state and [129, 5120] input, generated once on the GPU.""" + if seed not in _STATE_CACHE: + release_cache() # one seed's 18.75 GiB of expert weights, and its MoE, at a time + state = reference.generate_state(seed) + problems = reference.state_problems(state, seed) + assert not problems, problems + _STATE_CACHE[seed] = { + "state": state, + "input": reference.generate_input(seed), + "designed_state": None, + "moe": None, + "router": None, + } + return _STATE_CACHE[seed] + + +def case_input(seed: int, length: int) -> torch.Tensor: + return generated(seed)["input"][:length] + + +def designed_router() -> tuple[torch.Tensor, torch.Tensor]: + """fp32 router weight [256, 5120] and correction bias [256] for the designed rows. + + Noise of at most 2^-11 everywhere but the eight designed columns. Designed token t's + column holds (-1 + rank / 2) / 16 for each member of its set and -2 / 16 for every + other expert, so a row with 16 there and zero at every other designed column has the + designed logits plus noise of standard deviation about 0.02, far inside the 0.5 between + ranks. The bias, at most 2^-7, moves no selection. + """ + gen = torch.Generator(device="cuda") + gen.manual_seed(reference.tensor_seed(FIXTURE_SEED, "designed.gate.weight")) + weight = (2.0 * torch.rand(256, 5120, generator=gen, device="cuda") - 1.0) * 2.0**-11 + for columns, members in ((COLUMNS_A, SET_A), (COLUMNS_B, SET_B)): + for token, column in enumerate(columns): + weight[:, column] = OTHER_LOGIT / SPIKE + for expert_id, rank in zip(members, RANKS[token]): + weight[expert_id, column] = designed_logit(rank) / SPIKE + gen.manual_seed(reference.tensor_seed(FIXTURE_SEED, "designed.gate.e_score_correction_bias")) + bias = (2.0 * torch.rand(256, generator=gen, device="cuda") - 1.0) * 2.0**-7 + return weight.contiguous(), bias.contiguous() + + +def designed_rows(columns: tuple[int, ...], name: str) -> torch.Tensor: + """[4, 5120] bf16: unit normal noise, zero at every designed column but row t's own + ``columns[t]``, which holds 16.""" + gen = torch.Generator(device="cuda") + gen.manual_seed(reference.tensor_seed(FIXTURE_SEED, name)) + rows = torch.randn(len(columns), 5120, generator=gen, device="cuda") + rows[:, list(COLUMNS_A + COLUMNS_B)] = 0.0 + rows[torch.arange(len(columns)), list(columns)] = SPIKE + return rows.to(DTYPE) + + +def designed_ids(members: tuple[int, ...], token: int) -> list[int]: + """The members by descending designed logit: the routing the design forces on a token.""" + return [expert_id for _, expert_id in sorted(zip(RANKS[token], members), reverse=True)] + + +def designed_state(seed: int) -> "OrderedDict[str, torch.Tensor]": + """The seed's generated experts with the designed router.""" + entry = generated(seed) + if entry["designed_state"] is None: + weight, bias = designed_router() + router = {"gate.weight": weight, "gate.e_score_correction_bias": bias} + entry["designed_state"] = OrderedDict(entry["state"], **router) + return entry["designed_state"] + + +def loaded_moe(seed: int, router: str = "generated") -> BeamMoeBf16Sm103Tp1: + """The seed's cached MoE, loaded with the generated router or the "designed" one.""" + entry = generated(seed) + if entry["moe"] is None: + entry["moe"] = BeamMoeBf16Sm103Tp1(moe_config(), LAYER_IDX, device="cuda") + if entry["router"] != router: + entry["router"] = None # unknown until the load completes + entry["moe"].load_weights(entry["state"] if router == "generated" else designed_state(seed)) + entry["router"] = router + return entry["moe"] + + +def cheap_weights() -> "OrderedDict[str, torch.Tensor]": + """A valid 770-tensor mapping that allocates nothing: one expanded constant per dtype, + so a load that copied any of it would change the generated values it replaced.""" + constants = { + torch.float32: torch.full((), 0.5, dtype=torch.float32, device="cuda"), + torch.bfloat16: torch.full((), 0.25, dtype=torch.bfloat16, device="cuda"), + } + return OrderedDict( + (name, constants[dtype].expand(shape)) + for name, (shape, dtype) in reference.param_shapes().items() + ) + + +# -------------------------------------------------- op-path instrumentation + + +@contextlib.contextmanager +def recording(*wrappers, experts: bool = False): + """Log every call of these entries, in order, with its nesting depth, the arguments that + identify the path and its output; with ``experts``, also every expert the MoE asks for + and the rows it passes.""" + calls: list[dict] = [] + depth = [0] + saved = [] + for w in wrappers: + cls = type(w) + original = cls.__call__ + saved.append((cls, "__call__", original)) + + def logged(self, *args, _original=original, _name=cls.__name__, **kwargs): + tensors = [a for a in args if isinstance(a, torch.Tensor)] + entry = { + "op": _name, + "depth": depth[0], + "shapes": [list(t.shape) for t in tensors], + "data_ptrs": [t.data_ptr() for t in tensors], + } + entry.update({key: kwargs[key] for key in _RECORDED_KWARGS if key in kwargs}) + calls.append(entry) + depth[0] += 1 + try: + out = _original(self, *args, **kwargs) + finally: + depth[0] -= 1 + entry["output"] = out + return out + + cls.__call__ = logged + if experts: + original_expert = BeamMoeBf16Sm103Tp1._expert + saved.append((BeamMoeBf16Sm103Tp1, "_expert", original_expert)) + + def asked(self, expert_id, rows, _original=original_expert): + calls.append({"op": "expert", "depth": depth[0], "expert_id": expert_id, "rows": rows}) + return _original(self, expert_id, rows) + + BeamMoeBf16Sm103Tp1._expert = asked + try: + yield calls + finally: + for cls, name, original in reversed(saved): + setattr(cls, name, original) + + +def loggable(calls: list[dict]) -> list[dict]: + """The recorded calls without their tensors, for the evidence record.""" + keep = ("op", "depth", "shapes", "expert_id", *_RECORDED_KWARGS) + out = [] + for call in calls: + entry = {key: call[key] for key in keep if key in call} + if "rows" in call: + entry["rows"] = list(call["rows"].shape) + out.append(entry) + return out + + +@contextlib.contextmanager +def builtin_beam_forbidden(): + """The built-in Beam modeling must not run while the component does.""" + targets = [ + (builtin_beam.BeamMoE, "forward"), + (builtin_beam.BeamMoE, "combine"), + (builtin_beam.BeamRouter, "forward"), + (builtin_beam.BeamRouter, "route"), + (builtin_beam.BeamSwiGLU, "forward"), + (builtin_beam.BeamRMSNorm, "forward"), + (builtin_beam.BeamAttention, "forward"), + (builtin_beam.BeamDecoderLayer, "forward"), + ] + saved = [(cls, name, getattr(cls, name)) for cls, name in targets] + + def refuse(self, *args, **kwargs): + raise AssertionError(f"the built-in {type(self).__name__} ran inside the component") + + for cls, name, _ in saved: + setattr(cls, name, refuse) + try: + yield + finally: + for cls, name, original in reversed(saved): + setattr(cls, name, original) + + +@contextlib.contextmanager +def no_dispatch(): + """Fail on any op call: rejected inputs must be rejected before a kernel runs.""" + saved = [] + for w in OPS: + cls = type(w) + saved.append((cls, cls.__call__)) + + def refuse(self, *args, _name=cls.__name__, **kwargs): + raise AssertionError(f"{_name} dispatched before the input was rejected") + + cls.__call__ = refuse + try: + yield + finally: + for cls, original in reversed(saved): + cls.__call__ = original + + +# -------------------------------------------------- tests: configuration and construction + + +def test_configuration_is_the_fixed_scope() -> None: + """One configuration, no open choices: what the component accepts is what is tested.""" + assert CASES == ((17, 1), (17, 17), (17, 129), (41, 1), (41, 17), (41, 129)) + assert moe_module.SUPPORTED_LAYER_IDX == LAYER_IDX + for field, value in moe_module.SUPPORTED_CONFIG.items(): + assert MOE_CONFIG[field] == value, field + assert MOE_CONFIG["mlp_layer_types"][LAYER_IDX] == "sparse" + shapes = moe_module.weight_shapes(moe_config()) + assert list(shapes.items()) == list(reference.param_shapes().items()) + assert (reference.HIDDEN, reference.INTERMEDIATE, reference.NUM_EXPERTS) == ( + MOE_CONFIG["hidden_size"], + MOE_CONFIG["moe_intermediate_size"], + MOE_CONFIG["num_experts"], + ) + assert (reference.N_GROUP, reference.TOPK_GROUP, reference.TOP_K) == ( + MOE_CONFIG["n_group"], + MOE_CONFIG["topk_group"], + MOE_CONFIG["num_experts_per_tok"], + ) + assert ( + reference.ROUTED_SCALING_FACTOR + == MOE_CONFIG["routed_scaling_factor"] + == dispatch_combine.ROUTED_SCALING_FACTOR + ) + assert (reference.SEEDS, reference.LENGTHS) == (SEEDS, LENGTHS) + assert reference.MAX_TOKENS == max(LENGTHS) + assert torch.cuda.get_device_capability() == SM, "this suite certifies sm_103 only" + + +def test_moe_constructs_alone() -> None: + """The MoE is built directly from its configuration, without the model around it: a + router and 256 experts holding 770 tensors, every expert matrix its own storage.""" + + def refuse(self, *args, **kwargs): + raise AssertionError(f"{type(self).__name__} constructed") + + with pytest.MonkeyPatch.context() as mp: + for cls in ( + builtin_beam.BeamForCausalLM, + builtin_beam.BeamModel, + builtin_beam.BeamDecoderLayer, + builtin_beam.BeamMoE, + builtin_beam.BeamRouter, + builtin_beam.BeamSwiGLU, + ): + mp.setattr(cls, "__init__", refuse) + moe = BeamMoeBf16Sm103Tp1(moe_config(), LAYER_IDX, device="cuda") + assert isinstance(moe.gate, BeamMoeRouterBf16Sm103Tp1) + assert len(moe.experts) == 256 + assert all(isinstance(expert, BeamMoeExpertBf16Sm103Tp1) for expert in moe.experts) + assert [name for name, _ in moe.named_children()] == ["gate", "experts"] + tensors = moe.checkpoint_tensors() + assert list(tensors) == list(reference.param_shapes()) + for name, (shape, dtype) in reference.param_shapes().items(): + t = tensors[name] + assert tuple(t.shape) == shape and t.dtype == dtype and t.is_cuda, name + assert tensors["gate.weight"] is moe.gate.weight + assert tensors["gate.e_score_correction_bias"] is moe.gate.e_score_correction_bias + for expert_id in (0, 1, 31, 32, 128, 255): + expert = moe.experts[expert_id] + for attr in ("gate_proj", "up_proj", "down_proj"): + assert tensors[f"experts.{expert_id}.{attr}.weight"] is getattr(expert, attr) + pointers = [t.data_ptr() for name, t in tensors.items() if name.startswith("experts.")] + assert len(pointers) == 768 and len(set(pointers)) == 768 + assert "modeling_beam" not in Path(moe_module.__file__).read_text() + + +# -------------------------------------------------- tests: declared limits + +_BAD_CONFIGS = [ + ("moe_intermediate_size", dict(moe_intermediate_size=1536), LAYER_IDX, "moe_intermediate"), + ("routed_scaling_factor", dict(routed_scaling_factor=1.0), LAYER_IDX, "routed_scaling"), + ("hidden_act", dict(hidden_act="gelu"), LAYER_IDX, "unsupported hidden_act"), + ("num_experts", dict(num_experts=128), LAYER_IDX, "unsupported num_experts"), + ("dense_block_3", dict(mlp_layer_types=["dense"] * 4), LAYER_IDX, "sparse"), + ("block_2", {}, 2, "only the routers of sparse blocks"), +] + + +@pytest.mark.parametrize( + "case,overrides,layer_idx,match", _BAD_CONFIGS, ids=[c[0] for c in _BAD_CONFIGS] +) +def test_rejects_unsupported_configuration(case, overrides, layer_idx, match) -> None: + """Anything outside the one configuration is refused at construction.""" + with pytest.raises(ValueError, match=match): + BeamMoeBf16Sm103Tp1(moe_config(**overrides), layer_idx, device="cuda") + + +def test_rejects_a_host_device() -> None: + with pytest.raises(ValueError, match="sm_103"): + BeamMoeBf16Sm103Tp1(moe_config(), LAYER_IDX, device="cpu") + + +# -------------------------------------------------- tests: loading + + +def test_load_weights_maps_every_expert() -> None: + """Every checkpoint name lands in its own tensor, exactly: expert e's matrices in + ``experts[e]`` for all 256, the router's fp32 weight in ``gate.weight`` and the + selection bias in the router's ``e_score_correction_bias`` buffer.""" + seed = 17 + state = generated(seed)["state"] + moe = loaded_moe(seed) + loaded = moe.checkpoint_tensors() + assert list(loaded) == list(state) == list(reference.param_shapes()) + wrong = [] + for name, tensor in loaded.items(): + given = state[name] + if tensor.dtype != given.dtype or tensor.shape != given.shape: + wrong.append(f"{name}: {tuple(tensor.shape)} {tensor.dtype}") + elif not torch.equal(tensor, given): + wrong.append(f"{name}: values differ") + assert tensor.data_ptr() != given.data_ptr(), f"{name} aliases the given tensor" + assert not wrong, wrong[:8] + for expert_id, expert in enumerate(moe.experts): + for attr in ("gate_proj", "up_proj", "down_proj"): + given = state[f"experts.{expert_id}.{attr}.weight"] + assert torch.equal(getattr(expert, attr), given), (expert_id, attr) + router = moe.gate + assert router.weight.dtype == torch.float32 + assert torch.equal(router.weight, state["gate.weight"]) + assert [name for name, _ in router.named_buffers()] == ["e_score_correction_bias"] + assert [name for name, _ in router.named_parameters()] == ["weight"] + bias = router.e_score_correction_bias + assert bias.dtype == torch.float32 + assert torch.equal(bias, state["gate.e_score_correction_bias"]) + + +_BAD_LOADS = [ + ("last_down_fp32", TypeError, r"experts\.255\.down_proj\.weight must be torch\.bfloat16"), + ("last_gate_transposed", ValueError, r"experts\.255\.gate_proj\.weight must be \[2560, 5120\]"), + ("router_weight_bf16", TypeError, r"gate\.weight must be torch\.float32"), + ("missing_bias", ValueError, r"missing \['gate\.e_score_correction_bias'\]"), + ("unexpected_name", ValueError, r"unexpected \['experts\.256\.gate_proj\.weight'\]"), +] + + +@pytest.mark.parametrize("case,error,match", _BAD_LOADS, ids=[c for c, _, _ in _BAD_LOADS]) +def test_rejected_load_changes_nothing(case, error, match) -> None: + """Every tensor is checked before any is copied: a rejected load writes no parameter or + buffer, leaves every value as it was, and the MoE still runs.""" + seed = 17 + entry = generated(seed) + moe = loaded_moe(seed) + bad = cheap_weights() + if case == "last_down_fp32": + bad["experts.255.down_proj.weight"] = torch.zeros((), device="cuda").expand(5120, 2560) + elif case == "last_gate_transposed": + bad["experts.255.gate_proj.weight"] = bad["experts.255.gate_proj.weight"].t() + elif case == "router_weight_bf16": + bad["gate.weight"] = bad["gate.weight"].to(DTYPE) + elif case == "missing_bias": + del bad["gate.e_score_correction_bias"] + elif case == "unexpected_name": + bad["experts.256.gate_proj.weight"] = bad["experts.255.gate_proj.weight"] + names = ("gate.weight", "gate.e_score_correction_bias") + tuple( + f"experts.255.{n}" for n in reference.EXPERT_SHAPES + ) + snapshots = {name: moe.checkpoint_tensors()[name].clone() for name in names} + versions = tensor_versions(moe) + try: + with no_dispatch(), pytest.raises(error, match=match): + moe.load_weights(bad) + finally: + written = sorted(n for n, v in tensor_versions(moe).items() if v != versions[n]) + if written: + entry["router"] = None # reload the generated state before the MoE is used again + current = moe.checkpoint_tensors() + changed = [name for name, value in snapshots.items() if not torch.equal(current[name], value)] + assert not written, f"{len(written)} tensors written by a rejected load: {written[:8]}" + assert not changed, changed + out = moe(case_input(seed, 1)) + assert out.shape == (1, 5120) and bool(torch.isfinite(out).all()) + + +_BAD_INPUTS = [ + ("fp32_input", TypeError, "must be torch.bfloat16"), + ("fp16_input", TypeError, "must be torch.bfloat16"), + ("batched_3d", ValueError, r"must be \[tokens, 5120\]"), + ("width_5119", ValueError, r"must be \[tokens, 5120\]"), + ("zero_tokens", ValueError, "at least one token"), + ("host_input", ValueError, "MoE's device"), + ("not_a_tensor", TypeError, "must be a torch.Tensor"), + ("not_loaded", RuntimeError, "load_weights must be called"), +] + + +@pytest.mark.parametrize("case,error,match", _BAD_INPUTS, ids=[c for c, _, _ in _BAD_INPUTS]) +def test_rejects_unsupported_inputs_before_dispatch(case, error, match) -> None: + """Each declared restriction raises before any op runs and before any expert is asked + for.""" + if case == "not_loaded": + moe = BeamMoeBf16Sm103Tp1(moe_config(), LAYER_IDX, device="cuda") + x = torch.randn(17, 5120, device="cuda").to(DTYPE) + else: + moe = loaded_moe(17) + x = case_input(17, 17) + bad = { + "fp32_input": lambda: x.float(), + "fp16_input": lambda: x.half(), + "batched_3d": lambda: x.unsqueeze(0), + "width_5119": lambda: x[:, :5119], + "zero_tokens": lambda: x[:0], + "host_input": lambda: x.cpu(), + "not_a_tensor": lambda: x.tolist(), + "not_loaded": lambda: x, + }[case]() + with recording(experts=True) as asked, no_dispatch(), pytest.raises(error, match=match): + moe(bad) + assert asked == [], loggable(asked) + + +# -------------------------------------------------- tests: the six cases + + +@pytest.mark.parametrize("seed,length", CASES, ids=CASE_IDS) +def test_moe_matches_reference(seed, length, ieee_fp32) -> None: + """One case: the forward's own routing (the router op's mapping of that call) selects + exactly the reference's groups and experts in order, its fp32 steps and the output are + within their rules, and the call leaves its input and the parameters untouched.""" + moe = loaded_moe(seed) + state = generated(seed)["state"] + x = case_input(seed, length) + x_before = x.clone() + versions = tensor_versions(moe) + with recording(beam_grouped_topk_router) as calls: + out = moe(x) + torch.cuda.synchronize() + after = tensor_versions(moe) + (router_call,) = calls + route = router_call["output"] + ref_route = reference.route(x, state["gate.weight"], state["gate.e_score_correction_bias"]) + ref_out = reference.moe(x, state) + results = [output_contract(out, length)] + results.append(compare("moe_output", out, ref_out, "moe_output")) + results.extend(compare_routes(route, ref_route)) + results.append(weight_invariants(route)) + results.append(unchanged(x, x_before, versions, after)) + _write_evidence( + f"moe-seed{seed}-T{length}", + { + "test": "test_moe_matches_reference", + "case": {"seed": seed, "length": length}, + "rules": {k: [v[0], v[1], str(v[2])] for k, v in RULES.items()}, + "results": results, + "selected_experts": sorted(set(route["ids"].reshape(-1).tolist())), + "input_sha256": tensor_sha256(x), + "environment": _environment(ieee_fp32), + }, + ) + assert_all_pass(results) + + +# -------------------------------------------------- tests: the component path + + +def test_forward_runs_the_component_path(ieee_fp32) -> None: + """One case: the router op once, then BeamDispatchCombine, inside which every selected + expert -- ascending, once each -- receives exactly its tokens' rows in token order and + runs its own three projections and SwiGLU; every op's guard admits its arguments and no + built-in Beam module runs.""" + seed, length = 17, 17 + moe = loaded_moe(seed) + x = case_input(seed, length) + with builtin_beam_forbidden(), validating(*OPS), recording(*OPS, experts=True) as calls: + out = moe(x) + torch.cuda.synchronize() + _write_evidence( + f"path-seed{seed}-T{length}", + { + "test": "test_forward_runs_the_component_path", + "case": {"seed": seed, "length": length}, + "moe": f"{type(moe).__module__}.{type(moe).__qualname__}", + "op_calls": loggable(calls), + "environment": _environment(ieee_fp32), + }, + ) + + assert type(moe).__module__.endswith("moe_bf16__sm_103__tp1.moe") + assert [c["op"] for c in calls[:2]] == ["BeamGroupedTopkRouter", "BeamDispatchCombine"] + router, dispatch = calls[:2] + assert router["depth"] == dispatch["depth"] == 0 + assert {k: router[k] for k in ("n_group", "topk_group", "top_k", "eps")} == { + "n_group": 8, + "topk_group": 4, + "top_k": 8, + "eps": 1e-20, + } + assert router["shapes"] == [[length, 5120], [256, 5120], [256]] + assert dispatch["shapes"] == [[length, 5120], [length, 8], [length, 8]] + assert dispatch["routed_scaling_factor"] == 2.5 + assert dispatch["output"] is out + # Everything after the dispatch entry ran inside it. + nested = calls[2:] + assert all(c["depth"] == 1 for c in nested), loggable(nested) + ids = router["output"]["ids"] + selected = sorted(set(ids.reshape(-1).tolist())) + assert [c["op"] for c in nested] == list(EXPERT_CALLS) * len(selected), loggable(nested) + rows_total = 0 + for i, expert_id in enumerate(selected): + asked, gate, up, act, down = nested[len(EXPERT_CALLS) * i : len(EXPERT_CALLS) * (i + 1)] + assert asked["expert_id"] == expert_id + tokens = (ids == expert_id).any(dim=-1).nonzero().squeeze(-1) + n = int(tokens.numel()) + assert n >= 1 and bitwise_equal(asked["rows"], x[tokens]), expert_id + assert gate["shapes"] == up["shapes"] == [[n, 5120], [5120, 2560]], expert_id + assert act["shapes"] == [[n, 2560], [n, 2560]], expert_id + assert down["shapes"] == [[n, 2560], [2560, 5120]], expert_id + # Expert e's own stored matrices, read in place, on the rows it was given. + expert = moe.experts[expert_id] + assert gate["data_ptrs"] == [asked["rows"].data_ptr(), expert.gate_proj.data_ptr()] + assert up["data_ptrs"] == [asked["rows"].data_ptr(), expert.up_proj.data_ptr()] + assert down["data_ptrs"][1] == expert.down_proj.data_ptr() + rows_total += n + assert rows_total == 8 * length + assert out.shape == (length, 5120) and out.dtype == DTYPE and out.is_cuda + assert torch.cuda.get_device_capability(out.device) == SM + + +# -------------------------------------------------- tests: designed fixtures + + +def test_expert_identity_fixture(ieee_fp32, monkeypatch) -> None: + """Designed logits route four tokens to the same eight experts -- both ends of the id + range and of four groups -- in four different orders. Every one of the eight gets all + four rows, the other 248 are never asked for, the output matches the reference, and + putting expert 0 in expert 255's place fails the same comparison.""" + seed = 17 + moe = loaded_moe(seed, router="designed") + state = designed_state(seed) + x = designed_rows(COLUMNS_A, "designed.rows_a") + ref_route = reference.route(x, state["gate.weight"], state["gate.e_score_correction_bias"]) + ref_out = reference.moe(x, state) + with recording(beam_grouped_topk_router, experts=True) as calls: + out = moe(x) + route = calls[0]["output"] + asked = [c for c in calls if c["op"] == "expert"] + designed = [designed_ids(SET_A, token) for token in range(len(RANKS))] + runs = {"unchanged": compare("moe_output", out, ref_out, "moe_output")} + # experts[255] = experts[0], as ModuleList item assignment sets it, undone below. + monkeypatch.setattr(moe.experts, "255", moe.experts[0]) + assert moe.experts[255] is moe.experts[0] + runs["expert_255_is_expert_0"] = compare("moe_output", moe(x), ref_out, "moe_output") + monkeypatch.undo() + assert moe.experts[255] is not moe.experts[0] + runs["restored"] = compare("moe_output", moe(x), ref_out, "moe_output") + _write_evidence( + "fixture-expert-identity", + { + "test": "test_expert_identity_fixture", + "designed_ids": designed, + "ids": route["ids"].tolist(), + "weights": route["weights"].tolist(), + "expert_calls": loggable(asked), + "runs": runs, + "environment": _environment(ieee_fp32), + }, + ) + assert route["ids"].tolist() == designed, route["ids"].tolist() + assert ref_route["ids"].tolist() == designed, ref_route["ids"].tolist() + assert torch.equal(route["groups"], ref_route["groups"]) + assert sorted(route["groups"][0].tolist()) == [0, 1, 4, 7] + assert [c["expert_id"] for c in asked] == sorted(SET_A) + for call in asked: + assert bitwise_equal(call["rows"], x), call["expert_id"] + assert runs["unchanged"]["passed"] is True, runs["unchanged"]["problems"] + faulted = runs["expert_255_is_expert_0"] + assert faulted["passed"] is False and faulted["violations"] > 0, faulted + assert runs["restored"]["passed"] is True, runs["restored"]["problems"] + + +def test_changed_routing_between_calls(ieee_fp32) -> None: + """Rows routed to expert set A, then rows routed to the disjoint set B, then A again: + B's call routes as designed and asks for no expert of A, and the repeated A call + reproduces the first one bit for bit -- nothing of one call's routing survives into + the next.""" + seed = 17 + moe = loaded_moe(seed, router="designed") + state = designed_state(seed) + rows_a = designed_rows(COLUMNS_A, "designed.rows_a") + rows_b = designed_rows(COLUMNS_B, "designed.rows_b") + runs = [] + for x in (rows_a, rows_b, rows_a): + with recording(beam_grouped_topk_router, experts=True) as calls: + out = moe(x) + runs.append( + { + "out": out, + "ids": calls[0]["output"]["ids"], + "experts": [c["expert_id"] for c in calls if c["op"] == "expert"], + "vs_reference": compare("moe_output", out, reference.moe(x, state), "moe_output"), + } + ) + first, second, third = runs + _write_evidence( + "fixture-changed-routing", + { + "test": "test_changed_routing_between_calls", + "calls": [ + {"ids": r["ids"].tolist(), "experts": r["experts"], "result": r["vs_reference"]} + for r in runs + ], + "environment": _environment(ieee_fp32), + }, + ) + assert first["ids"].tolist() == [designed_ids(SET_A, t) for t in range(len(RANKS))] + assert second["ids"].tolist() == [designed_ids(SET_B, t) for t in range(len(RANKS))] + assert first["experts"] == third["experts"] == sorted(SET_A) + assert second["experts"] == sorted(SET_B) + assert not set(second["experts"]) & set(SET_A) + assert torch.equal(third["ids"], first["ids"]) + assert bitwise_equal(third["out"], first["out"]) + for run in runs: + assert run["vs_reference"]["passed"] is True, run["vs_reference"]["problems"] + + +def test_zero_input_gives_exact_zero() -> None: + """Zero hidden states: every score is 1/2, so every mixing weight is exactly 1/8, every + selected expert runs on zero rows, and the output is exactly zero -- no residual, bias + or other term reaches it.""" + moe = loaded_moe(17) + x = torch.zeros(17, 5120, dtype=DTYPE, device="cuda") + with recording(beam_grouped_topk_router, experts=True) as calls: + out = moe(x) + route = calls[0]["output"] + asked = [c for c in calls if c["op"] == "expert"] + assert out.shape == (17, 5120) and out.dtype == DTYPE + assert bool((out == 0).all()) + assert bool((route["scores"] == 0.5).all()) + assert bool((route["weights"] == 0.125).all()), route["weights"] + assert [c["expert_id"] for c in asked] == sorted(set(route["ids"].reshape(-1).tolist())) + + +def test_token_permutation_preserves_correspondence(ieee_fp32) -> None: + """Reversing the tokens reverses the routing and the output: each token keeps its own + experts and its own output row.""" + seed, length = 17, 17 + moe = loaded_moe(seed) + state = generated(seed)["state"] + x = case_input(seed, length) + perm = torch.arange(length - 1, -1, -1, device="cuda") + inverse = torch.argsort(perm) + outputs, routes = [], [] + for rows in (x, x[perm]): + with recording(beam_grouped_topk_router) as calls: + outputs.append(moe(rows)) + routes.append(calls[0]["output"]) + base, permuted = outputs + ref_out = reference.moe(x, state) + results = [ + compare_exact("ids_back_in_token_order", routes[1]["ids"][inverse], routes[0]["ids"]), + compare("output_back_in_token_order", permuted[inverse], base, "moe_output"), + compare("output_vs_reference", base, ref_out, "moe_output"), + compare("permuted_output_vs_reference", permuted, ref_out[perm], "moe_output"), + ] + _write_evidence( + f"permutation-seed{seed}-T{length}", + { + "test": "test_token_permutation_preserves_correspondence", + "results": results, + "environment": _environment(ieee_fp32), + }, + ) + assert_all_pass(results) + + +def test_detects_routed_scaling_faults(ieee_fp32, monkeypatch) -> None: + """The output comparison rejects the combination scaled by 1.0 or by 2.5 twice, and + accepts the unchanged and the restored MoE, on the same case and the same rule.""" + seed, length = 17, 17 + moe = loaded_moe(seed) + x = case_input(seed, length) + ref_out = reference.moe(x, generated(seed)["state"]) + runs = {"unchanged": compare("moe_output", moe(x), ref_out, "moe_output")} + for factor in (1.0, 6.25): + monkeypatch.setattr(dispatch_combine, "ROUTED_SCALING_FACTOR", factor) + runs[f"scaled_by_{factor}"] = compare("moe_output", moe(x), ref_out, "moe_output") + monkeypatch.undo() + runs["restored"] = compare("moe_output", moe(x), ref_out, "moe_output") + _write_evidence( + f"fault-scaling-seed{seed}-T{length}", + { + "test": "test_detects_routed_scaling_faults", + "runs": runs, + "environment": _environment(ieee_fp32), + }, + ) + assert dispatch_combine.ROUTED_SCALING_FACTOR == 2.5 + assert runs["unchanged"]["passed"] is True, runs["unchanged"]["problems"] + for factor in (1.0, 6.25): + faulted = runs[f"scaled_by_{factor}"] + assert faulted["passed"] is False and faulted["violations"] > 0, factor + assert runs["restored"]["passed"] is True, runs["restored"]["problems"] diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/_dispatch_combine_cases.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/_dispatch_combine_cases.py new file mode 100644 index 000000000000..f2328f547130 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/_dispatch_combine_cases.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Generated inputs, controlled experts and the plain-torch golden of the dispatch/combine tests. + +Block 3's MoE combination: hidden 5120, 256 experts, eight per token, routed scaling factor +2.5. The experts here are controlled functions rather than matrices: expert ``e`` maps a +BF16 row ``x`` to ``bf16(fp32(x) * A[e] + B[e])`` with fixed FP32 tables ``A`` and ``B`` +``[256, 5120]``, so every output depends on the expert's identity and on the gathered +token's values, and a row sent to the wrong expert or taken from the wrong token changes +it. Every tensor is generated on the host from a seed, so the tests need no files. +""" + +from __future__ import annotations + +import torch + +HIDDEN, EXPERTS, TOP_K, SCALE = 5120, 256, 8, 2.5 +SEEDS = (17, 41) +TOKENS = (1, 17, 129) +BF16, FP32, I64 = torch.bfloat16, torch.float32, torch.int64 +PROVIDER_SEED = 20261019 +PERM_SEED = 20261020 +FIXTURE_SEED = 20261021 + + +class AffineExperts: + """Expert ``e``: ``rows -> bf16(fp32(rows) * A[e] + B[e])``, row-wise; every call logged. + + ``calls`` holds ``(expert_id, rows, output)`` per call, in call order. An empty batch is + refused, so an expert called for no token fails the test that made the call. + """ + + def __init__(self, a: torch.Tensor, b: torch.Tensor): + self.a, self.b = a, b + self.calls: list[tuple[int, torch.Tensor, torch.Tensor]] = [] + + def __call__(self, expert_id: int, rows: torch.Tensor) -> torch.Tensor: + assert type(expert_id) is int, f"expert id must be a Python int, got {type(expert_id)}" + assert rows.shape[0] > 0, f"expert {expert_id} called with an empty batch" + out = (rows.float() * self.a[expert_id] + self.b[expert_id]).to(BF16) + self.calls.append((expert_id, rows, out)) + return out + + def fresh(self) -> "AffineExperts": + """The same experts with an empty call log.""" + return AffineExperts(self.a, self.b) + + +def affine_experts(device: str = "cuda", zero: bool = False) -> AffineExperts: + """Signed coefficients |A| in [0.5, 1.5) and offsets B ~ 0.25 N(0, 1); all zero if asked.""" + if zero: + return AffineExperts(*(torch.zeros(EXPERTS, HIDDEN, dtype=FP32, device=device),) * 2) + g = torch.Generator(device="cpu").manual_seed(PROVIDER_SEED) + magnitude = 0.5 + torch.rand(EXPERTS, HIDDEN, generator=g, dtype=FP32) + negative = torch.rand(EXPERTS, HIDDEN, generator=g, dtype=FP32) < 0.5 + a = torch.where(negative, -magnitude, magnitude) + b = 0.25 * torch.randn(EXPERTS, HIDDEN, generator=g, dtype=FP32) + return AffineExperts(a.to(device), b.to(device)) + + +def matrix_case(seed: int, tokens: int): + """``(hidden, ids, weights)`` on the host: 129 generated tokens, the first ``tokens`` used. + + Hidden states are unit-variance BF16. Each row's ids are eight distinct experts in + random (unsorted) order and its weights are normalized as the router normalizes them. + """ + g = torch.Generator(device="cpu").manual_seed(seed) + x = torch.randn(max(TOKENS), HIDDEN, generator=g, dtype=FP32).to(BF16) + ids = torch.stack([torch.randperm(EXPERTS, generator=g)[:TOP_K] for _ in range(max(TOKENS))]) + scores = 0.05 + torch.rand(max(TOKENS), TOP_K, generator=g, dtype=FP32) + w = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) + return x[:tokens].clone(), ids[:tokens].to(I64).clone(), w[:tokens].clone() + + +def permutation(seed: int, tokens: int) -> torch.Tensor: + g = torch.Generator(device="cpu").manual_seed(PERM_SEED + 1000 * seed + tokens) + return torch.randperm(tokens, generator=g) + + +SHARED_IDS = [ + [255, 3, 0, 128, 254, 1, 77, 200], + [0, 255, 9, 128, 31, 254, 2, 100], + [7, 0, 255, 1, 64, 129, 250, 33], + [254, 128, 0, 255, 5, 6, 11, 99], + [1, 2, 3, 4, 255, 0, 253, 252], + [128, 129, 130, 0, 131, 132, 255, 133], +] +WITNESS_IDS = [ + [12, 200, 57, 255, 0, 91, 143, 33], + [200, 12, 255, 64, 18, 91, 222, 5], + [0, 57, 143, 129, 200, 77, 255, 111], + [91, 33, 0, 12, 250, 129, 64, 7], +] +WITNESS_BASE_WEIGHTS = [0.34, 0.21, 0.14, 0.11, 0.08, 0.06, 0.04, 0.02] +ZERO_IDS = [ + [0, 17, 255, 42, 128, 3, 200, 91], + [255, 128, 7, 0, 64, 33, 250, 1], + [42, 0, 91, 254, 17, 255, 5, 129], +] +#: Tokens 0, 1 and 2 give experts 255, 33 and 129 a weight of exactly zero. +ZERO_ROUTE_SLOT = [2, 5, 7] + + +def fixtures() -> dict: + """The focused fixtures on the host, by name: ``(hidden, ids, weights, zero_experts)``. + + shared_extremes six distinct tokens all sharing experts 0 and 255 (and 128 / 254 / + 1 several times), unsorted slots, unequal weights that do not sum + to one, so a re-normalization would show + cancellation_witness signed routes of unequal size (weights 0.02 to 0.34, none a BF16 + value), so each rounding and the accumulation order are visible + zero_weight_route one route per token with weight exactly 0 + all_zero_weights every weight 0 + all_zero_outputs ordinary weights; the zero experts return exact zeros + """ + g = torch.Generator(device="cpu").manual_seed(FIXTURE_SEED) + fx = {} + x = torch.randn(6, HIDDEN, generator=g, dtype=FP32).to(BF16) + w = 0.05 + 0.35 * torch.rand(6, TOP_K, generator=g, dtype=FP32) + fx["shared_extremes"] = (x, torch.tensor(SHARED_IDS, dtype=I64), w, False) + x = torch.randn(4, HIDDEN, generator=g, dtype=FP32).to(BF16) + base = torch.tensor(WITNESS_BASE_WEIGHTS, dtype=FP32) + rows = [] + for _ in range(4): + perm = torch.randperm(TOP_K, generator=g) + jitter = 1.0 + 0.01 * torch.randn(TOP_K, generator=g, dtype=FP32) + rows.append(base[perm] * jitter) + raw = torch.stack(rows) + w = raw / (raw.sum(dim=-1, keepdim=True) + 1e-20) + fx["cancellation_witness"] = (x, torch.tensor(WITNESS_IDS, dtype=I64), w, False) + x = torch.randn(3, HIDDEN, generator=g, dtype=FP32).to(BF16) + raw = 0.05 + torch.rand(3, TOP_K, generator=g, dtype=FP32) + for t, j in enumerate(ZERO_ROUTE_SLOT): + raw[t, j] = 0.0 + w = raw / (raw.sum(dim=-1, keepdim=True) + 1e-20) + for t, j in enumerate(ZERO_ROUTE_SLOT): + w[t, j] = 0.0 + fx["zero_weight_route"] = (x, torch.tensor(ZERO_IDS, dtype=I64), w, False) + x = torch.randn(3, HIDDEN, generator=g, dtype=FP32).to(BF16) + fx["all_zero_weights"] = (x, torch.tensor(ZERO_IDS, dtype=I64), torch.zeros(3, TOP_K), False) + x = torch.randn(3, HIDDEN, generator=g, dtype=FP32).to(BF16) + raw = 0.05 + torch.rand(3, TOP_K, generator=g, dtype=FP32) + w = raw / (raw.sum(dim=-1, keepdim=True) + 1e-20) + fx["all_zero_outputs"] = (x, torch.tensor(ZERO_IDS, dtype=I64), w, True) + return fx + + +#: Wrong answers the golden can produce on request, one step changed each. +FAULTS = ( + "wrong_scatter_mapping", # contributions added to token (t + 1) mod T + "no_route_cast", # weighted route kept fp32 + "bf16_product", # weight rounded to bf16, product in bf16 + "cast_after_accumulate", # fp32 routes summed, the sum rounded to bf16 before scaling + "bf16_accumulator", # accumulator held in bf16 + "overwrite", # a contribution replaces the accumulator row instead of adding + "weights_slot_rotated", # each route takes the next slot's weight + "renormalized_weights", # weights re-normalized per token + "scale_omitted", + "scale_duplicated", +) + + +def golden(x, ids, w, experts, fault: str | None = None, skip_slots=None) -> dict: + """The source combination in plain torch, every step, on the inputs' device. + + Written independently of the component: every expert id 0..255 is enumerated, each + expert's (token, slot) pairs come from the routing table by ``nonzero``, and each + expert's rounded routes are added to their tokens' rows of an fp32 accumulator + (``acc[t] += route``; a token appears at most once per expert). Returns the stages + ``product`` / ``route`` (assignment rows in call order), ``acc`` (before scaling), + ``scaled`` and ``output``, plus ``calls`` = ``[(expert, tokens)]``. ``skip_slots[t]`` + drops slot ``skip_slots[t]`` of token ``t`` from the table. + """ + tokens = x.shape[0] + wk = w / (w.sum(dim=-1, keepdim=True) + 1e-20) if fault == "renormalized_weights" else w + acc = torch.zeros( + tokens, HIDDEN, dtype=BF16 if fault == "bf16_accumulator" else FP32, device=x.device + ) + stages = {"product": [], "route": [], "calls": []} + for e in range(EXPERTS): + token_slot = (ids == e).nonzero() + if skip_slots is not None: + keep = [slot != skip_slots[token] for token, slot in token_slot.tolist()] + token_slot = token_slot[torch.tensor(keep, dtype=torch.bool, device=ids.device)] + if token_slot.shape[0] == 0: + continue + t, j = token_slot[:, 0], token_slot[:, 1] + out = experts(e, x[t]) + jw = (j + 1) % TOP_K if fault == "weights_slot_rotated" else j + weight = wk[t, jw] + if fault == "bf16_product": + product = out * weight.to(BF16).unsqueeze(1) + route = product + else: + product = out.float() * weight.unsqueeze(1) + route = ( + product if fault in ("no_route_cast", "cast_after_accumulate") else product.to(BF16) + ) + dst = (t + 1) % tokens if fault == "wrong_scatter_mapping" else t + if fault == "overwrite": + acc[dst] = route.float() + elif fault == "bf16_accumulator": + acc[dst] = (acc[dst].float() + route.float()).to(BF16) + else: + acc[dst] += route.float() + stages["product"].append(product) + stages["route"].append(route) + stages["calls"].append((e, t)) + acc = acc.float() + if fault == "cast_after_accumulate": + acc = acc.to(BF16).float() + scaled = {"scale_omitted": acc, "scale_duplicated": acc * SCALE * SCALE}.get(fault, acc * SCALE) + stages.update(acc=acc, scaled=scaled, output=scaled.to(BF16)) + stages["product"] = torch.cat(stages["product"]) + stages["route"] = torch.cat(stages["route"]) + return stages + + +def bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + """Same dtype, shape and bits (so -0.0 differs from +0.0).""" + if a.dtype != b.dtype or a.shape != b.shape: + return False + view = {BF16: torch.int16, FP32: torch.int32}.get(a.dtype) + if view is None: + return bool(torch.equal(a, b.to(a.device))) + return bool(torch.equal(a.contiguous().view(view), b.to(a.device).contiguous().view(view))) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/moe/test_modeling_v2_beam_dispatch_combine.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/moe/test_modeling_v2_beam_dispatch_combine.py new file mode 100644 index 000000000000..bfc7abbd7235 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/moe/test_modeling_v2_beam_dispatch_combine.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_dispatch_combine entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_dispatch_combine_bf16__sm_103__tp1.ops.moe.beam_dispatch_combine import ( # noqa: E501 + beam_dispatch_combine_op as op, +) + +__extra_import_path__ = ["..", "../../../.."] +import _dispatch_combine_cases as cases # noqa: E402 +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_dispatch_combine requires a CUDA device" + + +@pytest.mark.parametrize("seed", cases.SEEDS) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:26] for c in op.CELLS]) +def test_certified_cells(cell, seed, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bit for bit: both sides form the same correctly rounded fp32 products, round each + route to bf16 once and add the routes in fp32 in the same expert order, so a single + differing bit is a different computation, not a rounding allowance.""" + certifying(op) + spec = cell.spec + x, ids, w = (t.cuda() for t in cases.matrix_case(seed, spec["tokens"])) + experts = cases.affine_experts() + assert (spec["hidden_size"], spec["num_experts"], spec["top_k"]) == ( + cases.HIDDEN, + cases.EXPERTS, + cases.TOP_K, + ) + called, referenced = experts.fresh(), experts.fresh() + out = op(x, ids, w, called, spec["routed_scaling_factor"]) + ref = op.reference(x, ids, w, referenced, spec["routed_scaling_factor"]) + assert out.dtype == torch.bfloat16 and out.shape == (spec["tokens"], cases.HIDDEN) + assert cases.bitwise_equal(out, ref) + # One call per selected expert, ascending, on both sides. + assert ( + [e for e, _, _ in called.calls] + == [e for e, _, _ in referenced.calls] + == sorted(set(ids.reshape(-1).tolist())) + ) + + +def test_reference_keeps_the_rounded_route() -> None: + """On the cancellation witness the reference differs from the same combination with + fp32 routes: rounding each weighted route to bf16 is the source's semantics.""" + x, ids, w, _ = ( + t.cuda() if isinstance(t, torch.Tensor) else t + for t in cases.fixtures()["cancellation_witness"] + ) + experts = cases.affine_experts() + ref = op.reference(x, ids, w, experts.fresh(), cases.SCALE) + gold = cases.golden(x, ids, w, experts.fresh()) + unrounded = cases.golden(x, ids, w, experts.fresh(), fault="no_route_cast") + assert cases.bitwise_equal(ref, gold["output"]) + assert not torch.equal(ref, unrounded["output"]) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/test_dispatch_combine.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/test_dispatch_combine.py new file mode 100644 index 000000000000..bc0a4902f544 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/test_dispatch_combine.py @@ -0,0 +1,344 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the Beam MoE dispatch and weighted combination of block 3: bf16 / sm_103 / tp1. + +Configuration under test, fixed in ``_dispatch_combine_cases.py``: hidden 5120, 256 +experts, eight per token, FP32 mixing weights, routed scaling factor 2.5; one sm_103 GPU; +eager; controlled affine experts instead of expert matrices; seeds {17, 41} x tokens +{1, 17, 129}, and focused fixtures. + +What is compared. Every call goes through the public ``beam_dispatch_combine``. Its output +must equal, bit for bit, the plain-torch golden of the source combination in +``_dispatch_combine_cases.py``: the steps are exact fp32 products, bf16 roundings and fp32 +additions in a fixed order, so any correct implementation reproduces them exactly. The +expert calls themselves -- which experts, in which order, with which rows -- are checked +against the routing table. + +Invocation:: + + python -m pytest tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1 + +Dependencies: a TensorRT-LLM build for sm_103. No model files. +""" + +from __future__ import annotations + +import inspect + +import pytest +import torch +from torch.overrides import TorchFunctionMode, resolve_name + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_dispatch_combine_bf16__sm_103__tp1 import ( # noqa: E501 + dispatch_combine, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_dispatch_combine_bf16__sm_103__tp1.dispatch_combine import ( # noqa: E501 + beam_dispatch_combine, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_dispatch_combine_bf16__sm_103__tp1.ops.moe.beam_dispatch_combine import ( # noqa: E501 + beam_dispatch_combine_op, +) + +__extra_import_path__ = ["."] +import _dispatch_combine_cases as cases # noqa: E402 + +assert torch.cuda.is_available(), "the Beam MoE dispatch/combine tests require a CUDA device" + +BF16, FP32, I64 = torch.bfloat16, torch.float32, torch.int64 + + +@pytest.fixture(scope="module") +def experts(): + return cases.affine_experts() + + +def call(x, ids, w, experts): + """One public call on fresh copies of the experts' log; returns (output, calls).""" + provider = experts.fresh() + out = beam_dispatch_combine(x, ids, w, provider) + return out, provider.calls + + +def check_calls(calls, x, ids): + """Exactly the selected experts, ascending, once each, each with its tokens' rows in + ascending token order: every one of the T*8 assignments once and nothing else.""" + selected = sorted(set(ids.reshape(-1).tolist())) + assert [e for e, _, _ in calls] == selected + assignments = 0 + for e, rows, _ in calls: + tokens = (ids == e).any(dim=1).nonzero().squeeze(1) + assert rows.shape[0] == tokens.numel() >= 1, e + assert cases.bitwise_equal(rows, x[tokens]), f"expert {e} received the wrong rows" + assignments += rows.shape[0] + assert assignments == ids.numel() + + +def run_checked(x, ids, w, experts, zero=False): + """Public call on the GPU with every output and immutability check; returns + (output, calls, golden).""" + x, ids, w = x.cuda(), ids.cuda(), w.cuda() + before = (x.clone(), ids.clone(), w.clone()) + out, calls = call(x, ids, w, experts) + assert out.shape == (x.shape[0], 5120) and out.dtype == BF16 and out.device == x.device + assert bool(torch.isfinite(out).all()) + for name, now, then in zip( + ("hidden_states", "expert_ids", "mixing_weights"), (x, ids, w), before + ): + assert cases.bitwise_equal(now, then), f"{name} was modified" + check_calls(calls, x, ids) + gold = cases.golden(x, ids, w, experts.fresh()) + assert [e for e, _, _ in calls] == [e for e, _ in gold["calls"]] + return out, calls, gold + + +# ----------------------------------------------------------------------------- the interface + + +def test_public_callable_and_provider_interface(): + params = list(inspect.signature(beam_dispatch_combine).parameters) + assert params == ["hidden_states", "expert_ids", "mixing_weights", "expert_provider"] + doc = " ".join(inspect.getdoc(dispatch_combine).split()) + for phrase in ( + "expert_provider(expert_id, rows)", + "owns dispatch and combination", + "owns the expert computation", + "[n, 5120]", + ): + assert phrase in doc, phrase + assert ( + dispatch_combine.HIDDEN_SIZE, + dispatch_combine.NUM_EXPERTS, + dispatch_combine.EXPERTS_PER_TOKEN, + dispatch_combine.ROUTED_SCALING_FACTOR, + ) == (5120, 256, 8, 2.5) + + +def test_component_holds_no_parameters_and_allocates_nothing_persistent(experts): + """No router parameter, expert matrix, module or configuration object exists to build: + the component is a function, its entry holds no state, and after a call only the + returned tensor remains allocated.""" + assert inspect.isfunction(beam_dispatch_combine) + assert vars(beam_dispatch_combine_op) == {} + assert not any( + isinstance(v, (torch.nn.Module, torch.Tensor)) for v in vars(dispatch_combine).values() + ) + x, ids, w = (t.cuda() for t in cases.matrix_case(17, 17)) + a, b = experts.a, experts.b + torch.cuda.synchronize() + before = torch.cuda.memory_stats()["requested_bytes.all.current"] + out = beam_dispatch_combine(x, ids, w, lambda e, rows: (rows.float() * a[e] + b[e]).to(BF16)) + torch.cuda.synchronize() + after = torch.cuda.memory_stats()["requested_bytes.all.current"] + assert after - before == out.numel() * out.element_size() + + +_T = 17 + + +@pytest.mark.parametrize( + "case,error,match", + [ + ("hidden_rank_1", ValueError, r"hidden_states must be \[tokens, 5120\]"), + ("hidden_rank_3", ValueError, r"hidden_states must be \[tokens, 5120\]"), + ("hidden_width_5119", ValueError, r"hidden_states must be \[tokens, 5120\]"), + ("zero_tokens", ValueError, "at least one token"), + ("hidden_fp32", TypeError, "hidden_states must be torch.bfloat16"), + ("hidden_fp16", TypeError, "hidden_states must be torch.bfloat16"), + ("hidden_not_a_tensor", TypeError, "hidden_states must be a torch.Tensor"), + ("hidden_on_host", ValueError, "sm_103 CUDA device"), + ("ids_rank_1", ValueError, r"expert_ids must be \[tokens, 8\]"), + ("ids_7_slots", ValueError, r"expert_ids must be \[tokens, 8\]"), + ("ids_extra_row", ValueError, r"expert_ids must be \[tokens, 8\]"), + ("ids_int32", TypeError, "expert_ids must be torch.int64"), + ("ids_float", TypeError, "expert_ids must be torch.int64"), + ("ids_on_host", ValueError, "expert_ids must be on hidden_states' device"), + ("weights_7_slots", ValueError, r"mixing_weights must be \[tokens, 8\]"), + ("weights_missing_row", ValueError, r"mixing_weights must be \[tokens, 8\]"), + ("weights_transposed", ValueError, r"mixing_weights must be \[tokens, 8\]"), + ("weights_bf16", TypeError, "mixing_weights must be torch.float32"), + ("weights_fp64", TypeError, "mixing_weights must be torch.float32"), + ("weights_on_host", ValueError, "mixing_weights must be on hidden_states' device"), + ("provider_not_callable", TypeError, "expert_provider must be callable"), + ], +) +def test_rejects_malformed_inputs_before_any_device_work(case, error, match, experts): + """Each argument is named in its error, nothing is computed and no provider is called; + the checks read metadata only, so they never synchronize with the device.""" + x, ids, w = (t.cuda() for t in cases.matrix_case(41, _T)) + provider = experts.fresh() + args = {"hidden_states": x, "expert_ids": ids, "mixing_weights": w, "expert_provider": provider} + change = { + "hidden_rank_1": ("hidden_states", x[0]), + "hidden_rank_3": ("hidden_states", x.unsqueeze(0)), + "hidden_width_5119": ("hidden_states", x[:, :5119]), + "zero_tokens": ("hidden_states", x[:0]), + "hidden_fp32": ("hidden_states", x.float()), + "hidden_fp16": ("hidden_states", x.half()), + "hidden_not_a_tensor": ("hidden_states", x.tolist()), + "hidden_on_host": ("hidden_states", x.cpu()), + "ids_rank_1": ("expert_ids", ids.reshape(-1)), + "ids_7_slots": ("expert_ids", ids[:, :7]), + "ids_extra_row": ("expert_ids", torch.cat([ids, ids[:1]])), + "ids_int32": ("expert_ids", ids.int()), + "ids_float": ("expert_ids", ids.float()), + "ids_on_host": ("expert_ids", ids.cpu()), + "weights_7_slots": ("mixing_weights", w[:, :7]), + "weights_missing_row": ("mixing_weights", w[:-1]), + "weights_transposed": ("mixing_weights", w.t()), + "weights_bf16": ("mixing_weights", w.to(BF16)), + "weights_fp64": ("mixing_weights", w.double()), + "weights_on_host": ("mixing_weights", w.cpu()), + "provider_not_callable": ("expert_provider", None), + }[case] + args[change[0]] = change[1] + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + with pytest.raises(error, match=match): + beam_dispatch_combine(**args) + finally: + torch.cuda.set_sync_debug_mode("default") + assert provider.calls == [] + + +@pytest.mark.parametrize( + "case,error,match", + [ + ("fp32", TypeError, r"expert_provider\(\d+\) must return torch.bfloat16"), + ("narrow", ValueError, r"expert_provider\(\d+\) must return \[\d+, 5120\]"), + ("extra_row", ValueError, r"expert_provider\(\d+\) must return \[\d+, 5120\]"), + ("host", ValueError, r"expert_provider\(\d+\) must return a dense tensor on cuda"), + ("not_a_tensor", TypeError, r"expert_provider\(\d+\) must return a torch.Tensor"), + ], +) +def test_rejects_provider_outputs_outside_the_contract(case, error, match, experts): + x, ids, w = (t.cuda() for t in cases.matrix_case(17, _T)) + good = experts.fresh() + bad = { + "fp32": lambda e, rows: good(e, rows).float(), + "narrow": lambda e, rows: good(e, rows)[:, :5119], + "extra_row": lambda e, rows: torch.cat([good(e, rows), good(e, rows[:1])]), + "host": lambda e, rows: good(e, rows).cpu(), + "not_a_tensor": lambda e, rows: good(e, rows).tolist(), + }[case] + with pytest.raises(error, match=match): + beam_dispatch_combine(x, ids, w, bad) + assert len(good.calls) in (1, 2) # refused at the first expert's output + + +# ----------------------------------------------------------------------------- the matrix + + +@pytest.mark.parametrize("tokens", cases.TOKENS) +@pytest.mark.parametrize("seed", cases.SEEDS) +def test_matrix_matches_the_source_combination(seed, tokens, experts): + out, calls, gold = run_checked(*cases.matrix_case(seed, tokens), experts) + assert cases.bitwise_equal(out, gold["output"]), f"seed {seed}, {tokens} tokens" + # Each expert's output went back to its own tokens: the scattered golden is different. + if tokens > 1: + wrong = cases.golden( + *(t.cuda() for t in cases.matrix_case(seed, tokens)), + experts.fresh(), + fault="wrong_scatter_mapping", + ) + assert not torch.equal(wrong["output"], out) + + +# ----------------------------------------------------------------------------- fixtures + + +def test_shared_extreme_experts_keep_tokens_and_experts_apart(experts): + x, ids, w, _ = cases.fixtures()["shared_extremes"] + out, calls, gold = run_checked(x, ids, w, experts) + assert cases.bitwise_equal(out, gold["output"]) + by_expert = {e: rows for e, rows, _ in calls} + for e in (0, 255): # every token selects both ends of the id range, in different slots + assert cases.bitwise_equal(by_expert[e], x.cuda()) + assert calls[0][0] == 0 and calls[-1][0] == 255 + # Leakage across tokens, across slots and a re-normalization each change the output. + for fault in ("wrong_scatter_mapping", "weights_slot_rotated", "renormalized_weights"): + wrong = cases.golden(x.cuda(), ids.cuda(), w.cuda(), experts.fresh(), fault=fault) + assert not torch.equal(wrong["output"], out), fault + + +def test_cancellation_witness_needs_every_rounding_step_in_place(experts): + """Signed routes of unequal size: the rounded route, the fp32 accumulation and the single + final scale are each visible in the output.""" + x, ids, w, _ = cases.fixtures()["cancellation_witness"] + out, calls, gold = run_checked(x, ids, w, experts) + assert cases.bitwise_equal(out, gold["output"]) + route = gold["route"].float() + assert int((route > 0).sum()) > 0 and int((route < 0).sum()) > 0 + assert bool((w.cuda() != w.cuda().to(BF16).float()).all()) # no weight is a bf16 value + for fault in cases.FAULTS: + if fault == "renormalized_weights": + continue # these weights are already normalized + wrong = cases.golden(x.cuda(), ids.cuda(), w.cuda(), experts.fresh(), fault=fault) + assert not torch.equal(wrong["output"], out), fault + + +def test_zero_weight_routes_contribute_exactly_zero(experts): + x, ids, w, _ = cases.fixtures()["zero_weight_route"] + out, calls, gold = run_checked(x, ids, w, experts) + assert cases.bitwise_equal(out, gold["output"]) + # A zero-weight route still calls its expert, and removing it changes nothing. + called = [e for e, _, _ in calls] + for t, j in enumerate(cases.ZERO_ROUTE_SLOT): + assert int(ids[t, j]) in called + without = cases.golden( + x.cuda(), ids.cuda(), w.cuda(), experts.fresh(), skip_slots=cases.ZERO_ROUTE_SLOT + ) + assert cases.bitwise_equal(out, without["output"]) + assert bool((out != 0).any(dim=1).all()) + + +@pytest.mark.parametrize("name", ["all_zero_weights", "all_zero_outputs"]) +def test_all_zero_weights_or_outputs_give_exact_zero(name, experts): + x, ids, w, zero = cases.fixtures()[name] + provider = cases.affine_experts(zero=True) if zero else experts + out, _, _ = run_checked(x, ids, w, provider) + assert bool((out == 0).all()) + # The hidden states are nonzero, so no residual or bias term reaches the output. + assert bool((x != 0).any()) + + +@pytest.mark.parametrize("seed", cases.SEEDS) +@pytest.mark.parametrize("tokens", [t for t in cases.TOKENS if t > 1]) +def test_row_permutation_keeps_each_token_its_output(seed, tokens, experts): + x, ids, w = (t.cuda() for t in cases.matrix_case(seed, tokens)) + perm = cases.permutation(seed, tokens).cuda() + base, _ = call(x, ids, w, experts) + permuted, _ = call(x[perm], ids[perm], w[perm], experts) + assert cases.bitwise_equal(permuted[torch.argsort(perm)], base) + + +# ----------------------------------------------------------------------------- execution + + +class TorchCallTrace(TorchFunctionMode): + """The name of every torch call made inside the block and the devices of its outputs.""" + + def __init__(self): + super().__init__() + self.calls: list[tuple[str, list[str]]] = [] + + def __torch_function__(self, func, types, args=(), kwargs=None): + out = func(*args, **(kwargs or {})) + outs = out if isinstance(out, (tuple, list)) else [out] + devices = [str(t.device) for t in outs if isinstance(t, torch.Tensor)] + self.calls.append((resolve_name(func) or repr(func), devices)) + return out + + +def test_token_arithmetic_runs_on_the_gpu(experts): + """Gathers, products, roundings, the accumulation and the scale all produce tensors on + the input's device; only the expert list itself comes to the host (``tolist``).""" + x, ids, w = (t.cuda() for t in cases.matrix_case(41, 129)) + with TorchCallTrace() as trace: + beam_dispatch_combine(x, ids, w, experts.fresh()) + names = [name for name, _ in trace.calls] + short = {name.rsplit(".", 1)[-1] for name in names} + assert {"__getitem__", "to", "index_add_"} <= short and {"__mul__", "mul"} & short, short + produced = [device for _, devices in trace.calls for device in devices] + assert produced and all(device == str(x.device) for device in produced), trace.calls + assert not any("reduce" in n or "all_gather" in n or "distributed" in n for n in names) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/_expert_cases.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/_expert_cases.py new file mode 100644 index 000000000000..76d0d76333d9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/_expert_cases.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Generated inputs, the plain-torch golden and the comparison rules of the expert tests. + +The expert is one SwiGLU expert of Beam decoder block 3: hidden 5120, intermediate 2560, +bias-free BF16 weights [out, in], BF16 projections, fp32 SiLU and product, one BF16 +rounding before the down projection. Every tensor here is generated on the host from a +seed or written down exactly, so the tests need no files. +""" + +from __future__ import annotations + +import contextlib +import math +import os + +import torch +import torch.nn.functional as F + +HIDDEN, INTER = 5120, 2560 +SEEDS = (17, 41) +TOKENS = (1, 17, 129) +STAGES = ("gate", "up", "silu", "product", "act", "y") +STAGE_DTYPES = { + "gate": torch.bfloat16, + "up": torch.bfloat16, + "silu": torch.float32, + "product": torch.float32, + "act": torch.bfloat16, + "y": torch.bfloat16, +} +BF16_UNIT = 2.0**-7 +FP32_FORMULA = 2.0**-16 + +#: (atol, rtol) per stage for the generated matrix, elementwise +#: ``|actual - expected| <= atol + rtol * |expected|``. rtol is one bf16 unit at any +#: magnitude; each atol is four times, rounded up, the largest excess over that measured +#: between valid implementations of this expert at these shapes (cuBLAS kernels, +#: fp32 and fp64 GEMMs, split-K orders, SiLU and product formulas). +MATRIX_RULES = { + "gate": (4e-5, BF16_UNIT), + "up": (4e-5, BF16_UNIT), + "silu": (0.02, BF16_UNIT), + "product": (0.009, BF16_UNIT), + "act": (0.04, BF16_UNIT), + "y": (0.008, BF16_UNIT), +} +#: Fixtures whose every projection output is one exact product, with values away from +#: any bf16 rounding boundary: projections, activation and output are exact, and SiLU and +#: product are fp32 formulas of exact inputs, within 2^-16 relative. +EXACT_RULES = { + "gate": (0.0, 0.0), + "up": (0.0, 0.0), + "silu": (0.0, FP32_FORMULA), + "product": (0.0, FP32_FORMULA), + "act": (0.0, 0.0), + "y": (0.0, 0.0), +} + +_TF32_OVERRIDES = ("TORCH_ALLOW_TF32_CUBLAS_OVERRIDE", "NVIDIA_TF32_OVERRIDE") + + +# ----------------------------------------------------------------------------- numerics + + +def numerical_state() -> dict: + m = torch.backends.cuda.matmul + return { + "float32_matmul_precision": torch.get_float32_matmul_precision(), + "cuda_matmul_fp32_precision": m.fp32_precision, + "cudnn_fp32_precision": torch.backends.cudnn.fp32_precision, + "cuda_matmul_allow_tf32": m.allow_tf32, + "cudnn_allow_tf32": torch.backends.cudnn.allow_tf32, + "bf16_reduced_precision_reduction": m.allow_bf16_reduced_precision_reduction, + "autocast_cuda": torch.is_autocast_enabled("cuda"), + } + + +_CONTROLLED = { + "float32_matmul_precision": "highest", + "cuda_matmul_fp32_precision": "ieee", + "cudnn_fp32_precision": "ieee", + "cuda_matmul_allow_tf32": False, + "cudnn_allow_tf32": False, + "bf16_reduced_precision_reduction": False, + "autocast_cuda": False, +} + + +def assert_numerical_controls() -> None: + for var in _TF32_OVERRIDES: + assert var not in os.environ, ( + f"{var} substitutes TF32 for fp32 GEMMs; run with env -u {var}" + ) + state = numerical_state() + assert state == _CONTROLLED, f"numerical controls not in force: {state}" + + +@contextlib.contextmanager +def numerical_controls(): + """IEEE fp32 matmul, bf16 GEMMs without reduced-precision reduction, autocast off. + + The golden's fp32 steps and its bf16 GEMMs follow these switches; the comparison rules + assume them. They are asserted on entry and exit and restored afterwards. + """ + m = torch.backends.cuda.matmul + saved = ( + torch.get_float32_matmul_precision(), + torch.backends.cudnn.allow_tf32, + m.fp32_precision, + torch.backends.cudnn.fp32_precision, + m.allow_bf16_reduced_precision_reduction, + ) + torch.set_float32_matmul_precision("highest") + m.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + m.fp32_precision = "ieee" + torch.backends.cudnn.fp32_precision = "ieee" + m.allow_bf16_reduced_precision_reduction = False + try: + with torch.autocast("cuda", enabled=False): + assert_numerical_controls() + yield + assert_numerical_controls() + finally: + # The legacy switches first: their setters also move the newer ones restored after. + torch.set_float32_matmul_precision(saved[0]) + torch.backends.cudnn.allow_tf32 = saved[1] + m.fp32_precision = saved[2] + torch.backends.cudnn.fp32_precision = saved[3] + m.allow_bf16_reduced_precision_reduction = saved[4] + + +def golden(x: torch.Tensor, gate_w: torch.Tensor, up_w: torch.Tensor, down_w: torch.Tensor): + """The source expert in plain torch, every step, on the inputs' device.""" + gate = F.linear(x, gate_w) + up = F.linear(x, up_w) + silu = F.silu(gate.float()) + product = silu * up.float() + act = product.to(torch.bfloat16) + return dict(gate=gate, up=up, silu=silu, product=product, act=act, y=F.linear(act, down_w)) + + +def violations(actual: torch.Tensor, expected: torch.Tensor, atol: float, rtol: float) -> dict: + """Elementwise rule check, both sides widened exactly to fp64; a nonfinite value fails.""" + a = actual.detach().double() + e = expected.detach().to(actual.device).double() + diff = (a - e).abs() + ok = diff <= atol + rtol * e.abs() # NaN compares False: a nonfinite element fails + bad = ~ok + worst = int(torch.argmax(torch.where(ok, torch.zeros_like(diff), diff.nan_to_num(math.inf)))) + return { + "failing": int(bad.sum()), + "nonfinite": int((~torch.isfinite(a)).sum()), + "max_abs": float(diff.nan_to_num(math.inf).max()) if diff.numel() else 0.0, + "worst": (worst, float(a.reshape(-1)[worst]), float(e.reshape(-1)[worst])) + if diff.numel() + else None, + } + + +def assert_stages(actual: dict, expected: dict, rules: dict, what: str) -> None: + """Every stage: its dtype and shape, and each element within that stage's rule.""" + problems = [] + for stage in STAGES: + a, e = actual[stage], expected[stage] + if a.dtype != STAGE_DTYPES[stage]: + problems.append(f"{stage}: dtype {a.dtype}, expected {STAGE_DTYPES[stage]}") + if a.shape != e.shape: + problems.append(f"{stage}: shape {tuple(a.shape)}, expected {tuple(e.shape)}") + continue + found = violations(a, e, *rules[stage]) + if found["failing"] or found["nonfinite"]: + problems.append(f"{stage}: {found} under (atol, rtol) {rules[stage]}") + assert not problems, f"{what}: " + "; ".join(problems) + + +# ----------------------------------------------------------------------------- generated cases + + +def matrix_case(seed: int) -> tuple[dict, torch.Tensor]: + """Seed -> (weights by checkpoint name, the [129, 5120] input; case T uses its first T rows). + + Draw order gate, up, down, input, standard normal fp32 from a host generator; gate and + up scaled by 1/sqrt(5120), down by 1/sqrt(2560), each rounded once to bf16. + """ + g = torch.Generator(device="cpu").manual_seed(seed) + + def draw(rows: int, cols: int, scale: float) -> torch.Tensor: + return (torch.randn(rows, cols, generator=g, dtype=torch.float32) * scale).to( + torch.bfloat16 + ) + + weights = { + "gate_proj.weight": draw(INTER, HIDDEN, 1.0 / math.sqrt(HIDDEN)), + "up_proj.weight": draw(INTER, HIDDEN, 1.0 / math.sqrt(HIDDEN)), + "down_proj.weight": draw(HIDDEN, INTER, 1.0 / math.sqrt(INTER)), + } + return weights, draw(max(TOKENS), HIDDEN, 1.0) + + +def as_weights(gate: torch.Tensor, up: torch.Tensor, down: torch.Tensor) -> dict: + return {"gate_proj.weight": gate, "up_proj.weight": up, "down_proj.weight": down} + + +# ----------------------------------------------------------------------------- exact fixtures + + +def silu64(t: torch.Tensor) -> torch.Tensor: + return F.silu(t.double()) + + +def bf16_margin(v: torch.Tensor) -> torch.Tensor: + """Relative distance of each fp64 value from its nearest bf16 rounding boundary.""" + v = v.double() + nearest = v.to(torch.bfloat16) + bits = nearest.view(torch.int16) + above = (bits + 1).view(torch.bfloat16).double() # one unit further from zero + below = (bits - 1).view(torch.bfloat16).double() # one unit nearer to zero + n = nearest.double() + dist = torch.minimum((v - (n + above) / 2).abs(), (v - (n + below) / 2).abs()) + return torch.where(v == 0, torch.full_like(v, math.inf), dist / v.abs()) + + +def bf16_grid(lo_exp: int, hi_exp: int, powers_of_two: bool = True) -> torch.Tensor: + """Every positive bf16 value in [2^lo_exp, 2^(hi_exp + 1)), as fp64.""" + mantissas = torch.arange(0 if powers_of_two else 1, 128, dtype=torch.float64) + return torch.cat([(1 + mantissas / 128) * 2.0**e for e in range(lo_exp, hi_exp + 1)]) + + +def identity_block_weights() -> dict: + """gate[i] = x[i], up[i] = x[2560 + i], y[5119 - 2i] = act[i]: one exact term each.""" + i = torch.arange(INTER) + gate = torch.zeros(INTER, HIDDEN, dtype=torch.bfloat16) + up = torch.zeros(INTER, HIDDEN, dtype=torch.bfloat16) + down = torch.zeros(HIDDEN, INTER, dtype=torch.bfloat16) + gate[i, i] = 1 + up[i, i + INTER] = 1 + down[HIDDEN - 1 - 2 * i, i] = 1 + return as_weights(gate, up, down) + + +def identity_block_expected(gate: torch.Tensor, up: torch.Tensor) -> dict: + """What every stage must be, exactly, for exact bf16 `gate`/`up` through identity blocks.""" + s = silu64(gate) + act = (s.float().double() * up).float().to(torch.bfloat16) + y = torch.zeros(gate.shape[0], HIDDEN, dtype=torch.bfloat16) + y[:, HIDDEN - 1 - 2 * torch.arange(INTER)] = act + return dict( + gate=gate.to(torch.bfloat16), + up=up.to(torch.bfloat16), + silu=s.float(), + product=(s.float().double() * up).float(), + act=act, + y=y, + ) + + +def signed_gate_fixture() -> tuple[dict, torch.Tensor, dict]: + """Negative, zero and positive gates over four decades, power-of-two signed ups. + + Gates whose SiLU lies within 2^-16 of a bf16 rounding boundary are dropped, so every + stage is exact; with a power-of-two up the product only rescales the SiLU. + """ + magnitudes = [1000.0, 60.0, 24.0, 10.0, 6.0, 4.5, 3.0, 2.25, 1.75, 1.28125, 1.0, 0.875] + magnitudes += [0.625, 0.375, 0.25, 0.125, 2.0**-4, 2.0**-6, 2.0**-10, 2.0**-20, 2.0**-60] + candidates = torch.tensor( + [-m for m in magnitudes] + [0.0, -0.0] + magnitudes, dtype=torch.float64 + ) + gates = candidates[bf16_margin(silu64(candidates)) > FP32_FORMULA] + ups = torch.tensor( + [1.0, -1.0, 2.0, -2.0, 0.5, -0.5, 4.0, -0.25, 0.125, -8.0], dtype=torch.float64 + ) + i = torch.arange(INTER) + gate = torch.stack([gates[(i + 37 * t) % len(gates)] for t in range(4)]) + up = torch.stack([ups[(5 * i + 3 * t) % len(ups)] for t in range(4)]) + x = torch.cat([gate, up], dim=1).to(torch.bfloat16) + return identity_block_weights(), x, identity_block_expected(gate, up) + + +def precision_fixture() -> tuple[dict, torch.Tensor, dict, torch.Tensor]: + """Gate/up pairs where the fp32 SiLU and product, rounded once, differ from bf16 steps. + + Over every bf16 gate with |g| in [2^-8, 8) and every non-power-of-two bf16 up with |u| + in [0.5, 4), keep the pairs whose correct activation and whose prematurely rounded one + (bf16 SiLU, bf16 product) are both 2^-16 or more from a bf16 rounding boundary and + differ, then draw 4 x 2560 of them with a fixed seed. Returns the weights, the input, + the exact stages and the premature activation every element must differ from. + """ + g_pos, u_pos = bf16_grid(-8, 2), bf16_grid(-1, 1, powers_of_two=False) + g = torch.cat([-g_pos.flip(0), g_pos]) + u = torch.cat([-u_pos.flip(0), u_pos]) + s = silu64(g) + product = s.float().double()[:, None] * u[None, :] + premature = s.float().to(torch.bfloat16).double()[:, None] * u[None, :] + robust = (bf16_margin(product) > FP32_FORMULA) & (bf16_margin(s) > FP32_FORMULA)[:, None] + robust &= bf16_margin(premature) > FP32_FORMULA + differs = product.float().to(torch.bfloat16) != premature.to(torch.bfloat16) + rows, cols = torch.nonzero(robust & differs, as_tuple=True) + pick = torch.randperm(rows.numel(), generator=torch.Generator().manual_seed(20261020)) + pick = pick[: 4 * INTER] + gate, up = g[rows[pick]].reshape(4, INTER), u[cols[pick]].reshape(4, INTER) + x = torch.cat([gate, up], dim=1).to(torch.bfloat16) + premature_act = (s.float().to(torch.bfloat16).double()[rows[pick]] * u[cols[pick]]).to( + torch.bfloat16 + ) + return ( + identity_block_weights(), + x, + identity_block_expected(gate, up), + premature_act.reshape(4, INTER), + ) + + +def sparse_mapping_fixture() -> tuple[dict, torch.Tensor, dict]: + """Asymmetric one-term maps with power-of-two markers: + gate[i] = m_g(i) x[a(i)], up[i] = m_u(i) x[b(i)], y[d(i)] = m_d(i) act[i]. + + a and b together cover all 5120 input features, d is a permutation into the 5120 + outputs, and gate and up read different inputs with different markers, so a transposed + weight, swapped gate/up roles or a wrong down mapping all move some output. + """ + i = torch.arange(INTER) + a, b, d = (3 * i + 7) % HIDDEN, (3 * i + 2567) % HIDDEN, (HIDDEN - 1 - 3 * i) % HIDDEN + m_g = 2.0 ** ((i % 5) - 2).double() + m_u = -(2.0 ** ((i % 3) - 1).double()) + m_d = 2.0 ** ((i % 4) - 1).double() + candidates = torch.cat([-bf16_grid(-3, 0).flip(0), bf16_grid(-3, 0)]) + robust = torch.ones_like(candidates, dtype=torch.bool) + for marker in (0.25, 0.5, 1.0, 2.0, 4.0): + robust &= bf16_margin(silu64(candidates * marker)) > FP32_FORMULA + palette = candidates[robust] + palette = palette[torch.randperm(palette.numel(), generator=torch.Generator().manual_seed(7))] + x = torch.zeros(3, HIDDEN, dtype=torch.float64) + for t in range(3): + x[t, a] = palette[(7 * i + 13 * t) % 64] + sign = torch.where((i + 2 * t) % 3 == 0, -1.0, 1.0).double() + x[t, b] = sign * 2.0 ** (((i + t) % 5) - 2).double() + gate_w = torch.zeros(INTER, HIDDEN, dtype=torch.bfloat16) + up_w = torch.zeros(INTER, HIDDEN, dtype=torch.bfloat16) + down_w = torch.zeros(HIDDEN, INTER, dtype=torch.bfloat16) + gate_w[i, a] = m_g.to(torch.bfloat16) + up_w[i, b] = m_u.to(torch.bfloat16) + down_w[d, i] = m_d.to(torch.bfloat16) + gate, up = x[:, a] * m_g, x[:, b] * m_u + expected = identity_block_expected(gate, up) + y = torch.zeros(3, HIDDEN, dtype=torch.bfloat16) + y[:, d] = (expected["act"].double() * m_d).to(torch.bfloat16) + expected["y"] = y + return as_weights(gate_w, up_w, down_w), x.to(torch.bfloat16), expected diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/activation/test_modeling_v2_swiglu_fp32.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/activation/test_modeling_v2_swiglu_fp32.py new file mode 100644 index 000000000000..104a4b68c0f9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/activation/test_modeling_v2_swiglu_fp32.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the swiglu_fp32 entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_expert_bf16__sm_103__tp1.ops.activation.swiglu_fp32 import ( # noqa: E501 + swiglu_fp32 as op, +) + +__extra_import_path__ = ["..", "../../../.."] +import _expert_cases as cases # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "swiglu_fp32 requires a CUDA device" + + +@pytest.fixture(scope="module", autouse=True) +def controls(): + with cases.numerical_controls(): + yield + + +def _build(spec, seed): + """Gate and up as the expert's projections produce them: unit-variance bf16.""" + g = torch.Generator(device="cpu").manual_seed(seed) + shape = (spec["tokens"], spec["intermediate"]) + gate = torch.randn(shape, generator=g).to(spec["dtype"]).cuda() + up = torch.randn(shape, generator=g).to(spec["dtype"]).cuda() + return gate, up + + +@pytest.mark.parametrize("seed", cases.SEEDS) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:26] for c in op.CELLS]) +def test_certified_cells(cell, seed, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """SiLU and product within 2^-16 relative of the fp64 steps (fp32 formula spread on the + same inputs), the activation within one bf16 unit (a product that close to a rounding + boundary may round either way).""" + certifying(op) + gate, up = _build(cell.spec, seed) + out = op(gate, up) + ref = op.reference(gate, up) + assert {k: v.dtype for k, v in out.items()} == { + "silu": torch.float32, + "product": torch.float32, + "act": torch.bfloat16, + } + for name, rtol in (("silu", cases.FP32_FORMULA), ("product", cases.FP32_FORMULA)): + torch.testing.assert_close(out[name], ref[name], rtol=rtol, atol=0.0) + torch.testing.assert_close(out["act"], ref["act"], rtol=cases.BF16_UNIT, atol=0.0) + + +def test_bf16_steps_are_a_different_answer() -> None: + """On the precision fixture's pairs every activation is exact here, and every one + differs from the activation of a bf16 SiLU times up computed in bf16.""" + _, x, expected, premature_act = cases.precision_fixture() + gate, up = x[:, : cases.INTER].cuda(), x[:, cases.INTER :].cuda() + out = op(gate, up) + assert torch.equal(out["act"].cpu(), expected["act"]) + bf16_steps = torch.nn.functional.silu(gate) * up + assert torch.equal(bf16_steps.cpu(), premature_act) + assert bool((out["act"] != bf16_steps).all()) + + +@pytest.mark.parametrize( + "case,match", + [ + ("broadcast_up", "must have one shape"), + ("fp32_gate", "must be torch.bfloat16"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + gate, up = _build(op.CELLS[1].spec, seed=17) + if case == "broadcast_up": + up = up[:1] + else: + gate = gate.float() + with pytest.raises(AssertionError, match=match): + with validating(op): + op(gate, up) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/test_expert.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/test_expert.py new file mode 100644 index 000000000000..327144e19c9c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/test_expert.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for one Beam MoE routed expert of block 3: bf16 / sm_103 / tp1. + +Configuration under test, fixed in ``_expert_cases.py``: hidden 5120, intermediate 2560, +bias-free BF16 weights [out, in], BF16 already-normalized input; one sm_103 GPU; eager; +generated parameters; seeds {17, 41} x tokens {1, 17, 129}, and exact fixtures. + +What is compared. Every call goes through the component's ``forward``. The three +``cublas_mm`` results and the activation op's mapping of that same call are captured, and +every step -- gate, up, fp32 SiLU, fp32 product, bf16 activation, output -- is compared +with the source expert written in plain torch under ``_expert_cases.MATRIX_RULES``, or +exactly on the fixtures built so that every step is known exactly. + +Invocation (fp32 steps need the TF32 overrides removed before torch is imported):: + + 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 + +Dependencies: a TensorRT-LLM build for sm_103. No model files. +""" + +from __future__ import annotations + +import contextlib +import gc + +import pytest +import torch +from torch.overrides import TorchFunctionMode, resolve_name + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import CublasMm, cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_expert_bf16__sm_103__tp1.expert import ( # noqa: E501 + WEIGHT_SHAPES, + BeamMoeExpertBf16Sm103Tp1, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_expert_bf16__sm_103__tp1.ops.activation.swiglu_fp32 import ( # noqa: E501 + SwigluFp32, + swiglu_fp32, +) + +__extra_import_path__ = [".", "../../.."] +import _expert_cases as cases # noqa: E402 +from _validating import validating # noqa: E402 + +assert torch.cuda.is_available(), "the Beam MoE expert tests require a CUDA device" + +BF16 = torch.bfloat16 +_WEIGHT_BYTES = sum(2 * rows * cols for rows, cols in WEIGHT_SHAPES.values()) + + +@pytest.fixture(scope="module", autouse=True) +def controls(): + """IEEE fp32, no reduced-precision bf16 reduction and no autocast for the whole module.""" + with cases.numerical_controls(): + yield + + +# ----------------------------------------------------------------------------- observing a call + + +@contextlib.contextmanager +def observed_calls(): + """Every cublas_mm result and activation mapping made inside the block, unchanged.""" + seen = {"mm": [], "act": []} + mm_call, act_call = CublasMm.__call__, SwigluFp32.__call__ + + def mm(self, *args, **kwargs): + out = mm_call(self, *args, **kwargs) + seen["mm"].append((args, out)) + return out + + def act(self, *args, **kwargs): + out = act_call(self, *args, **kwargs) + seen["act"].append((args, out)) + return out + + CublasMm.__call__, SwigluFp32.__call__ = mm, act + try: + yield seen + finally: + CublasMm.__call__, SwigluFp32.__call__ = mm_call, act_call + + +def run(expert: BeamMoeExpertBf16Sm103Tp1, x: torch.Tensor) -> dict: + """Every step of one forward call, by name; ``y`` is what forward returned.""" + with validating(cublas_mm, swiglu_fp32), observed_calls() as seen: + y = expert(x) + assert len(seen["mm"]) == 3 and len(seen["act"]) == 1, seen + (gate_args, gate), (up_args, up), (down_args, out) = seen["mm"] + ((act_args, steps),) = seen["act"] + for args, weight in ( + (gate_args, expert.gate_proj), + (up_args, expert.up_proj), + (down_args, expert.down_proj), + ): + # The stored [out, in] weight itself, read column-major: no copy, no transpose kernel. + operand = args[1] + assert operand.data_ptr() == weight.data_ptr() and operand.stride() == (1, weight.shape[1]) + assert act_args[0] is gate and act_args[1] is up and down_args[0] is steps["act"] + assert out is y + return dict(gate=gate, up=up, **{k: steps[k] for k in ("silu", "product", "act")}, y=y) + + +class TorchCallTrace(TorchFunctionMode): + """The name of every torch call made inside the block and the devices of its outputs.""" + + def __init__(self): + super().__init__() + self.calls: list[tuple[str, list[str]]] = [] + + def __torch_function__(self, func, types, args=(), kwargs=None): + out = func(*args, **(kwargs or {})) + outs = out if isinstance(out, (tuple, list)) else [out] + devices = [str(t.device) for t in outs if isinstance(t, torch.Tensor)] + self.calls.append((resolve_name(func) or repr(func), devices)) + return out + + +def loaded(weights: dict) -> BeamMoeExpertBf16Sm103Tp1: + expert = BeamMoeExpertBf16Sm103Tp1(device="cuda") + expert.load_weights(weights) + return expert + + +def bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + return ( + a.dtype == b.dtype + and a.shape == b.shape + and torch.equal(a.view(torch.int16), b.view(torch.int16)) + ) + + +# ----------------------------------------------------------------------------- state + + +def _requested_cuda_bytes() -> int: + """Bytes the live tensors asked for; unlike memory_allocated, free of block rounding.""" + return torch.cuda.memory_stats()["requested_bytes.all.current"] + + +def test_construction_holds_one_expert_and_nothing_else(): + gc.collect() + torch.cuda.synchronize() + before = _requested_cuda_bytes() + expert = BeamMoeExpertBf16Sm103Tp1(device="cuda") + torch.cuda.synchronize() + assert _requested_cuda_bytes() - before == _WEIGHT_BYTES == 78_643_200 + inventory = { + name: (tuple(p.shape), p.dtype, p.device.type, p.requires_grad) + for name, p in expert.named_parameters() + } + assert inventory == { + "gate_proj": ((2560, 5120), BF16, "cuda", False), + "up_proj": ((2560, 5120), BF16, "cuda", False), + "down_proj": ((5120, 2560), BF16, "cuda", False), + } + assert sum(p.numel() for p in expert.parameters()) == 39_321_600 + assert list(expert.buffers()) == [] and list(expert.children()) == [] + assert list(expert.checkpoint_tensors()) == list(WEIGHT_SHAPES) + + +def test_construction_rejects_a_host_device(): + with pytest.raises(ValueError, match="sm_103 CUDA devices"): + BeamMoeExpertBf16Sm103Tp1(device="cpu") + + +@pytest.mark.parametrize("source_device", ["cpu", "cuda"]) +def test_load_copies_exact_values(source_device): + weights = {k: v.to(source_device) for k, v in cases.matrix_case(17)[0].items()} + expert = loaded(weights) + for name, tensor in expert.checkpoint_tensors().items(): + assert tensor.device.type == "cuda" and tensor.dtype == BF16 + assert tuple(tensor.shape) == WEIGHT_SHAPES[name] + assert bitwise_equal(tensor.cpu(), weights[name].cpu()), name + assert tensor.data_ptr() != weights[name].data_ptr() + + +@pytest.mark.parametrize( + "case,error,match", + [ + ("not_loaded", RuntimeError, "load_weights must be called"), + ("rank_1", ValueError, r"must be \[tokens, 5120\]"), + ("rank_3", ValueError, r"must be \[tokens, 5120\]"), + ("width_5119", ValueError, r"must be \[tokens, 5120\]"), + ("zero_tokens", ValueError, "at least one token"), + ("fp32", TypeError, "must be torch.bfloat16"), + ("fp16", TypeError, "must be torch.bfloat16"), + ("host", ValueError, "expert's device"), + ("not_a_tensor", TypeError, "must be a torch.Tensor"), + ], +) +def test_rejects_unsupported_inputs_before_any_kernel(case, error, match): + weights, x_all = cases.matrix_case(17) + expert = BeamMoeExpertBf16Sm103Tp1(device="cuda") + if case != "not_loaded": + expert.load_weights(weights) + x = x_all[:17].cuda() + bad = { + "not_loaded": x, + "rank_1": x[0], + "rank_3": x.unsqueeze(0), + "width_5119": x[:, :5119], + "zero_tokens": x[:0], + "fp32": x.float(), + "fp16": x.half(), + "host": x.cpu(), + "not_a_tensor": x.tolist(), + }[case] + with observed_calls() as seen, pytest.raises(error, match=match): + expert(bad) + assert seen["mm"] == [] and seen["act"] == [] + + +@pytest.mark.parametrize( + "case,error,match", + [ + ("down_fp32", TypeError, "down_proj.weight must be torch.bfloat16"), + ("down_transposed", ValueError, r"down_proj.weight must be \[5120, 2560\]"), + ("gate_fp16", TypeError, "gate_proj.weight must be torch.bfloat16"), + ("up_transposed", ValueError, r"up_proj.weight must be \[2560, 5120\]"), + ("missing_down", ValueError, r"missing \['down_proj.weight'\]"), + ("unexpected_bias", ValueError, r"unexpected \['down_proj.bias'\]"), + ("down_not_a_tensor", TypeError, "down_proj.weight must be a torch.Tensor"), + ("down_meta", ValueError, "dense host or CUDA tensor"), + ("not_a_mapping", TypeError, "weights must be a mapping"), + ], +) +def test_rejected_load_leaves_the_expert_unchanged(case, error, match): + """Every other tensor is a valid replacement; a bad down tensor is the last one checked.""" + weights, x_all = cases.matrix_case(17) + expert = loaded(weights) + x = x_all[:17].cuda() + before = {k: v.clone() for k, v in expert.checkpoint_tensors().items()} + y_before = expert(x) + valid = cases.matrix_case(41)[0] + gate, up, down = valid.values() + replacement = { + "down_fp32": dict(valid, **{"down_proj.weight": down.float()}), + "down_transposed": dict(valid, **{"down_proj.weight": down.reshape(2560, 5120)}), + "gate_fp16": dict(valid, **{"gate_proj.weight": gate.half()}), + "up_transposed": dict(valid, **{"up_proj.weight": up.reshape(5120, 2560)}), + "missing_down": {k: v for k, v in valid.items() if k != "down_proj.weight"}, + "unexpected_bias": dict(valid, **{"down_proj.bias": down[:, 0]}), + "down_not_a_tensor": dict(valid, **{"down_proj.weight": down.float().numpy()}), + "down_meta": dict(valid, **{"down_proj.weight": torch.empty_like(down, device="meta")}), + "not_a_mapping": list(valid.items()), + }[case] + with pytest.raises(error, match=match): + expert.load_weights(replacement) + for name, tensor in expert.checkpoint_tensors().items(): + assert bitwise_equal(tensor, before[name]), name + assert bitwise_equal(expert(x), y_before) + + +def test_rejected_first_load_leaves_the_expert_unloaded(): + weights = dict(cases.matrix_case(17)[0]) + weights["down_proj.weight"] = weights["down_proj.weight"].float() + expert = BeamMoeExpertBf16Sm103Tp1(device="cuda") + with pytest.raises(TypeError): + expert.load_weights(weights) + with pytest.raises(RuntimeError, match="load_weights must be called"): + expert(torch.zeros(1, 5120, dtype=BF16, device="cuda")) + + +# ----------------------------------------------------------------------------- the matrix + + +@pytest.mark.parametrize("tokens", cases.TOKENS) +@pytest.mark.parametrize("seed", cases.SEEDS) +def test_matrix_every_step_matches_the_source_expert(seed, tokens): + weights, x_all = cases.matrix_case(seed) + expert = loaded(weights) + x = x_all[:tokens].cuda() + x_before = x.clone() + params_before = {k: v.clone() for k, v in expert.checkpoint_tensors().items()} + + steps = run(expert, x) + y = steps["y"] + assert y.shape == (tokens, 5120) and y.dtype == BF16 and y.device == x.device + assert bool(torch.isfinite(y).all()) + assert bitwise_equal(x, x_before) + for name, tensor in expert.checkpoint_tensors().items(): + assert bitwise_equal(tensor, params_before[name]), name + # Observing the call does not change it. + assert bitwise_equal(expert(x), y) + + gold = cases.golden(x, *(t.cuda() for t in weights.values())) + cases.assert_stages(steps, gold, cases.MATRIX_RULES, f"seed {seed}, {tokens} tokens") + + +# ----------------------------------------------------------------------------- fixtures + + +def test_zero_input_gives_exact_zero(): + expert = loaded(cases.matrix_case(17)[0]) + steps = run(expert, torch.zeros(17, 5120, dtype=BF16, device="cuda")) + for stage, tensor in steps.items(): + assert tensor.dtype == cases.STAGE_DTYPES[stage] + assert bool((tensor == 0).all()), stage + + +def test_signed_gates_exercise_silu_exactly(): + weights, x, expected = cases.signed_gate_fixture() + gates = expected["gate"].double() + assert int((gates < 0).sum()) >= 2000 and int((gates > 0).sum()) >= 2000 + assert int((gates == 0).sum()) > 0 and int((gates.abs() >= 24).sum()) > 0 + steps = run(loaded(weights), x.cuda()) + cases.assert_stages(steps, expected, cases.EXACT_RULES, "signed gates") + # SiLU is not linear on these gates: it is neither the gate nor half of it. + silu = steps["silu"].double().cpu() + assert not torch.equal(silu, gates) and not torch.equal(silu, gates / 2) + + +def test_precision_fixture_tells_fp32_activation_from_bf16_steps(): + weights, x, expected, premature_act = cases.precision_fixture() + steps = run(loaded(weights), x.cuda()) + cases.assert_stages(steps, expected, cases.EXACT_RULES, "precision fixture") + # Every element of the exact activation differs from the twice-rounded one. + assert bool((steps["act"].cpu().double() != premature_act.double()).all()) + + +def test_sparse_mapping_fixes_orientation_and_feature_maps(): + weights, x, expected = cases.sparse_mapping_fixture() + steps = run(loaded(weights), x.cuda()) + cases.assert_stages(steps, expected, cases.EXACT_RULES, "sparse mapping") + # The fixture separates the maps it is meant to check: each wrong reading moves y. + gate_w, up_w, down_w = (t.cuda() for t in weights.values()) + xd = x.cuda() + wrong = { + "gate/up swapped": cases.golden(xd, up_w, gate_w, down_w), + "gate read transposed": cases.golden(xd, gate_w.reshape(5120, 2560).t(), up_w, down_w), + "down read transposed": cases.golden(xd, gate_w, up_w, down_w.reshape(2560, 5120).t()), + } + for name, gold in wrong.items(): + assert not torch.equal(gold["y"].cpu(), expected["y"]), name + + +def test_row_permutation_keeps_tokens_apart(): + weights, x_all = cases.matrix_case(17) + expert = loaded(weights) + x = x_all.cuda() + perm = torch.randperm(x.shape[0], generator=torch.Generator().manual_seed(20261019)).cuda() + inverse = torch.argsort(perm) + base = run(expert, x) + permuted = run(expert, x[perm]) + back = {stage: t[inverse] for stage, t in permuted.items()} + cases.assert_stages(back, base, cases.MATRIX_RULES, "row permutation") + # A token-mixing expert would fail the same rule. + mixed = torch.roll(permuted["y"], 1, dims=0)[inverse] + assert cases.violations(mixed, base["y"], *cases.MATRIX_RULES["y"])["failing"] > 0 + + +def test_no_residual_weighting_or_collective(): + """Zero down weights with a nonzero input: an added residual, bias or any other term + would leave the output nonzero. Exact outputs on the fixtures above rule out a mixing + weight or the routed scaling factor.""" + weights = dict(cases.matrix_case(17)[0]) + weights["down_proj.weight"] = torch.zeros_like(weights["down_proj.weight"]) + expert = loaded(weights) + x = cases.matrix_case(17)[1][:17].cuda() + with TorchCallTrace() as trace: + steps = run(expert, x) + assert bool((steps["y"] == 0).all()) + assert all(bool((steps[s] != 0).any()) for s in ("gate", "up", "act")) + names = [name for name, _ in trace.calls] + assert not any("reduce" in n or "all_gather" in n or "distributed" in n for n in names), names + + +def test_forward_runs_on_the_gpu(): + weights, x_all = cases.matrix_case(41) + expert = loaded(weights) + x = x_all[:129].cuda() + with TorchCallTrace() as trace: + expert(x) + names = [name for name, _ in trace.calls] + assert sum("cublas_mm" in name for name in names) == 3, names + assert "torch.nn.functional.silu" in names, names + produced = [device for _, devices in trace.calls for device in devices] + assert produced and all(device == str(x.device) for device in produced), trace.calls diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/_router_cases.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/_router_cases.py new file mode 100644 index 000000000000..6ef460f0d26e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/_router_cases.py @@ -0,0 +1,420 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs and comparison rules shared by the Beam MoE router's tests. + +Everything here is generated on the CPU from fixed seeds, so the bytes a test routes are +the same on every machine and can be checked against routes captured elsewhere. + +Matrix cases, per seed: a CPU generator seeded with it draws the FP32 router weight +``randn(256, 5120) / sqrt(5120)``, then the FP32 correction bias ``0.05 * randn(256)``, +then FP32 token rows ``randn(T, 5120)``, scaled to unit RMS and rounded to BF16 -- the +already-normalized input the router receives. Weight and bias are drawn first, so they +are the same for every token count of a seed. + +Controlled fixtures fix the logits exactly instead: token row ``r`` is one-hot at column +``r`` and the seed-17 weight's column ``r`` is replaced by the logits that row should +produce, so ``logits[r] = weight[:, r]`` whatever the GEMM's reduction order. Each fixture +is a decision the router has to get right and says which. +""" + +from __future__ import annotations + +import math + +import torch + +HIDDEN, EXPERTS, GROUPS, PER_GROUP, TOPK_GROUP, TOP_K = 5120, 256, 8, 32, 4, 8 +NORM_EPS = 1e-20 +SEEDS = (17, 41) +TOKENS = (1, 17, 129) +CASES = tuple((seed, tokens) for seed in SEEDS for tokens in TOKENS) +CASE_IDS = [f"seed{seed}-T{tokens}" for seed, tokens in CASES] + +#: Elementwise |actual - reference| <= atol (rtol 0) per compared tensor, both widened to +#: fp64. Each atol is four times the largest difference measured between valid fp32 +#: evaluations of this routing at these shapes and inputs -- GEMM reduction orders and +#: layouts, sigmoid forms, normalization forms -- rounded up to one significant digit. +#: A route is never excused by a tolerance: groups and ids must match exactly, in order. +RULES = { + "logits": 8e-6, + "scores": 3e-6, + "selection_scores": 3e-6, + "group_scores": 2e-6, + "weights": 8e-7, +} + +#: |sum(weights) - 1| bound for every row whose selected scores are ordinary: eight fp32 +#: terms summed, divided and summed again err by at most 15 units of 2^-24 (8.9e-7). +NORM_ATOL = 1e-6 + + +# ----------------------------------------------------------------------------- matrix + + +def matrix_case(seed: int, tokens: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """(weight [256, 5120] fp32, correction bias [256] fp32, hidden [T, 5120] bf16) on CPU.""" + g = torch.Generator(device="cpu").manual_seed(seed) + weight = torch.randn(EXPERTS, HIDDEN, generator=g, dtype=torch.float32) * ( + 1.0 / math.sqrt(HIDDEN) + ) + bias = torch.randn(EXPERTS, generator=g, dtype=torch.float32) * 0.05 + x = torch.randn(tokens, HIDDEN, generator=g, dtype=torch.float32) + x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True)) + return weight, bias, x.to(torch.bfloat16) + + +# ----------------------------------------------------------------------------- controlled fixtures + + +def _logit(p: float) -> float: + return math.log(p / (1.0 - p)) + + +def _row(default: float) -> list: + return [default] * EXPERTS + + +def _set_group(row: list, group: int, values) -> None: + for j, v in enumerate(values): + row[group * PER_GROUP + j] = v + + +def _logits_for_scores(rows) -> torch.Tensor: + """FP32 logits whose sigmoid is each row's intended score.""" + return torch.tensor([[_logit(p) for p in r] for r in rows], dtype=torch.float64).float() + + +def _bias_fixtures() -> list: + # Each group has two lead experts; the rest sit at distinct values just above 0.30. + p = [0.30 + 0.0001 * ((e * 37) % EXPERTS) for e in range(EXPERTS)] + leads = { + 0: (0.70, 0.69), + 1: (0.68, 0.67), + 2: (0.66, 0.65), + 3: (0.64, 0.63), + 4: (0.62, 0.61), + 5: (0.60, 0.59), + 6: (0.50, 0.49), + 7: (0.48, 0.47), + } + for group, (a, b) in leads.items(): + p[group * PER_GROUP], p[group * PER_GROUP + 1] = a, b + logits = _logits_for_scores([p]) + shifted = torch.zeros(EXPERTS, dtype=torch.float32) + shifted[6 * PER_GROUP] = 0.25 + shifted[6 * PER_GROUP + 1] = 0.27 + return [ + dict( + name="bias_zero", + rows=("lead_experts",), + logits=logits, + bias=torch.zeros(EXPERTS), + intent="no bias: groups [0, 1, 2, 3], experts 0, 1, 32, 33, 64, 65, 96, 97", + ), + dict( + name="bias_shift", + rows=("lead_experts",), + logits=logits, + bias=shifted, + intent="+0.25/+0.27 on experts 192/193 lifts group 6 to first: groups [6, 0, 1, 2]; " + "the weights stay the normalized unbiased scores of the selected experts", + ), + ] + + +def _group_boundary() -> dict: + rows = [] + for second in (0.62, 0.67): + r = _row(0.20) + _set_group(r, 0, [0.90, 0.89]) + _set_group(r, 1, [0.88, 0.87]) + _set_group(r, 2, [0.86, 0.85]) + _set_group(r, 3, [0.655, 0.645]) # incumbent: top-2 sum 1.300 + _set_group(r, 4, [0.66, second]) # challenger: 1.28 below, 1.33 above + _set_group(r, 5, [0.95] + [0.05] * (PER_GROUP - 1)) # best single expert, sum 1.00 + _set_group(r, 6, [0.30, 0.29]) + _set_group(r, 7, [0.28, 0.27]) + rows.append(r) + return dict( + name="group_boundary", + rows=("challenger_below", "challenger_above"), + logits=_logits_for_scores(rows), + bias=torch.zeros(EXPERTS), + intent="only group 4's second expert moves (0.62 -> 0.67): groups [0, 1, 2, 3] -> " + "[0, 1, 2, 4]; scoring groups by their best expert alone would pick group 5", + ) + + +def _expert_boundary() -> dict: + rows = [] + for mover in (0.79, 0.82): + r = _row(0.30) + _set_group(r, 0, [0.95, 0.93, mover]) # third in its group: group scores do not move + _set_group(r, 1, [0.91, 0.89]) + _set_group(r, 2, [0.87, 0.85]) + _set_group(r, 3, [0.83, 0.81]) + _set_group(r, 4, [0.60, 0.58]) + _set_group(r, 5, [0.59, 0.57]) + _set_group(r, 6, [0.56, 0.55]) + _set_group(r, 7, [0.54, 0.53]) + rows.append(r) + return dict( + name="expert_boundary", + rows=("mover_below", "mover_above"), + logits=_logits_for_scores(rows), + bias=torch.zeros(EXPERTS), + intent="expert 2 moves 0.79 -> 0.82 across the 8th/9th boundary and expert 97 (0.81) " + "drops out; the selected groups do not change", + ) + + +EXCLUSION_SPIKE = 7 * PER_GROUP + 5 + + +def _exclusion() -> dict: + logits = [0.0005 * ((e * 37) % EXPERTS) for e in range(EXPERTS)] + logits[EXCLUSION_SPIKE] = 6.0 + bias = [] + for e in range(EXPERTS): + group = e // PER_GROUP + if group <= 3: + bias.append(-2.0) + elif group <= 6: + bias.append(-5.0) + else: + bias.append(0.0 if e == EXCLUSION_SPIKE else -10.0) + return dict( + name="exclusion", + rows=("spike_in_excluded_group",), + logits=torch.tensor([logits], dtype=torch.float32), + bias=torch.tensor(bias, dtype=torch.float32), + intent=f"expert {EXCLUSION_SPIKE} has the highest selection score but its group loses; " + "every retained selection score is negative, so skipping the exclusion or masking with " + "zero instead of -inf both select an excluded expert", + ) + + +def _ties() -> dict: + def top7() -> list: + r = _row(0.30) + _set_group(r, 0, [0.95, 0.94]) + _set_group(r, 1, [0.93, 0.92]) + _set_group(r, 2, [0.91, 0.90]) + _set_group(r, 3, [0.89, 0.60]) + for g in range(4, 8): + _set_group(r, g, [0.50 - 0.01 * g, 0.49 - 0.01 * g]) + return r + + rows, names = [], [] + r = top7() + r[1 * PER_GROUP + 7] = 0.70 + r[1 * PER_GROUP + 25] = 0.70 + rows.append(r) + names.append("expert_boundary_tie_same_group_39_57") + r = top7() + r[0 * PER_GROUP + 3] = 0.70 + r[3 * PER_GROUP + 30] = 0.70 + rows.append(r) + names.append("expert_boundary_tie_distant_3_126") + r = _row(0.30) + _set_group(r, 0, [0.95, 0.94]) + r[1 * PER_GROUP + 1] = 0.92 + r[1 * PER_GROUP + 30] = 0.92 + _set_group(r, 2, [0.91, 0.90]) + _set_group(r, 3, [0.89, 0.88]) + for g in range(4, 8): + _set_group(r, g, [0.50 - 0.01 * g, 0.49 - 0.01 * g]) + rows.append(r) + names.append("expert_order_tie_33_62") + r = _row(0.30) + _set_group(r, 0, [0.95, 0.94]) + _set_group(r, 1, [0.93, 0.92]) + _set_group(r, 2, [0.91, 0.90]) + _set_group(r, 3, [0.80, 0.78]) + _set_group(r, 6, [0.80, 0.78]) + _set_group(r, 4, [0.50, 0.49]) + _set_group(r, 5, [0.48, 0.47]) + _set_group(r, 7, [0.46, 0.45]) + rows.append(r) + names.append("group_boundary_tie_3_6") + r = _row(0.30) + _set_group(r, 0, [0.95, 0.94]) + _set_group(r, 1, [0.90, 0.89]) + _set_group(r, 4, [0.90, 0.89]) + _set_group(r, 2, [0.85, 0.84]) + _set_group(r, 3, [0.50, 0.49]) + _set_group(r, 5, [0.48, 0.47]) + _set_group(r, 6, [0.46, 0.45]) + _set_group(r, 7, [0.44, 0.43]) + rows.append(r) + names.append("group_order_tie_1_4") + rows.append(_row(0.5)) + names.append("all_equal") + r = _row(0.30) + _set_group(r, 0, [0.95, 0.94]) + _set_group(r, 1, [0.93, 0.92]) + _set_group(r, 2, [0.91, 0.90]) + _set_group(r, 3, [0.89]) + for g, js in ((0, (5, 17, 31)), (1, (2, 28)), (2, (9, 10, 20)), (3, (4, 30))): + for j in js: + r[g * PER_GROUP + j] = 0.70 + for g in range(4, 8): + _set_group(r, g, [0.50 - 0.01 * g, 0.49 - 0.01 * g]) + rows.append(r) + names.append("expert_boundary_tie_10way") + return dict( + name="ties", + rows=tuple(names), + logits=_logits_for_scores(rows), + bias=torch.zeros(EXPERTS), + intent="bitwise-equal scores at the 8th/9th expert boundary, inside the top-8 order, at " + "the 4th/5th group boundary, inside the top-4 group order, everywhere, and a 10-way " + "boundary tie", + ) + + +#: The order `torch.topk(..., sorted=True)` returns for each tie row -- the source router's +#: order, measured on sm_103. A tie at a selection boundary keeps the lower index; tied +#: values inside the selection are not in index order. +TIE_ROUTES = { + "expert_boundary_tie_same_group_39_57": ([0, 1, 2, 3], [0, 1, 32, 33, 64, 65, 96, 39]), + "expert_boundary_tie_distant_3_126": ([0, 1, 2, 3], [0, 1, 32, 33, 64, 65, 96, 3]), + "expert_order_tie_33_62": ([0, 1, 2, 3], [0, 1, 33, 62, 64, 65, 96, 97]), + "group_boundary_tie_3_6": ([0, 1, 2, 3], [0, 1, 32, 33, 64, 65, 96, 97]), + "group_order_tie_1_4": ([0, 4, 1, 2], [0, 1, 128, 32, 33, 129, 64, 65]), + "all_equal": ([1, 0, 2, 3], [7, 6, 4, 5, 1, 0, 2, 3]), + "expert_boundary_tie_10way": ([0, 1, 2, 3], [0, 1, 32, 33, 64, 65, 96, 5]), +} + +EPSILON_IDS = [0, 1, 32, 33, 64, 65, 96, 97] + + +def _epsilon() -> dict: + def bias_of(e: int) -> float: + group, j = divmod(e, PER_GROUP) + return 0.9 - 0.1 * group - 0.001 * j if j < 2 else -0.5 - 0.01 * group - 0.0001 * j + + zero_row = [-120.0] * EXPERTS + near_row = [-60.0] * EXPERTS + for k, e in enumerate(EPSILON_IDS): + near_row[e] = -47.5 - 0.15 * k + return dict( + name="epsilon", + rows=("zero_sum", "near_epsilon_sum"), + logits=torch.tensor([zero_row, near_row], dtype=torch.float32), + bias=torch.tensor([bias_of(e) for e in range(EXPERTS)], dtype=torch.float32), + intent="the bias fixes the selection to experts 0, 1, 32, 33, 64, 65, 96, 97; row 0's " + "scores underflow to 0 (weights 0, not NaN), row 1's selected scores sum to ~1e-20, so " + "its weights sum visibly below 1 with exactly the 1e-20 denominator term", + ) + + +def fixtures() -> dict: + """Every controlled fixture by name, in a fixed order.""" + specs = [*_bias_fixtures(), _group_boundary(), _expert_boundary(), _exclusion(), _ties()] + specs.append(_epsilon()) + return {spec["name"]: spec for spec in specs} + + +def fixture_tensors(spec: dict) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """(weight, correction bias, hidden) on CPU that make the router see `spec`'s logits.""" + logits = spec["logits"].float() + rows = logits.shape[0] + weight = matrix_case(17, 1)[0].clone() + weight[:, :rows] = logits.t() + hidden = torch.zeros(rows, HIDDEN, dtype=torch.bfloat16) + hidden[torch.arange(rows), torch.arange(rows)] = 1.0 + return weight.contiguous(), spec["bias"].float().contiguous(), hidden + + +# ----------------------------------------------------------------------------- comparison + + +def compare(name: str, actual: torch.Tensor, expected: torch.Tensor, atol: float) -> dict: + """Score one fp32 tensor against `atol` (rtol 0); non-finite values must match exactly.""" + result = {"tensor": name, "atol": atol, "rtol": 0.0, "shape": list(expected.shape)} + problems = [] + if tuple(actual.shape) != tuple(expected.shape): + problems.append(f"shape {tuple(actual.shape)} != {tuple(expected.shape)}") + if actual.dtype != torch.float32 or expected.dtype != torch.float32: + problems.append(f"dtypes {actual.dtype}/{expected.dtype}, need torch.float32") + if not problems: + a, e = actual.detach().double(), expected.detach().double() + finite_a, finite_e = torch.isfinite(a), torch.isfinite(e) + nonfinite_mismatch = int(((finite_a != finite_e) | (~finite_e & (a != e))).sum()) + both = finite_a & finite_e + diff = (a[both] - e[both]).abs() + violations = int((diff > atol).sum()) + result.update( + max_abs=float(diff.max()) if diff.numel() else 0.0, + unequal=int((diff > 0).sum()), + violations=violations, + nonfinite_mismatch=nonfinite_mismatch, + numel=int(e.numel()), + ) + if violations: + problems.append(f"{violations} elements exceed atol {atol:g} (max {diff.max():.3e})") + if nonfinite_mismatch: + problems.append(f"{nonfinite_mismatch} non-finite positions differ") + result["problems"] = problems + result["passed"] = not problems + return result + + +def compare_exact(name: str, actual: torch.Tensor, expected: torch.Tensor) -> dict: + """Ordered integer ids: equal element for element, never sorted first.""" + same = tuple(actual.shape) == tuple(expected.shape) and actual.dtype == expected.dtype + rows = ( + (actual == expected).all(dim=-1) + if same + else torch.zeros(expected.shape[0], dtype=torch.bool, device=expected.device) + ) + mismatched = [int(i) for i in torch.nonzero(~rows).flatten().cpu()] + return { + "tensor": name, + "rule": "exact ordered equality", + "shape": list(expected.shape), + "dtype": str(actual.dtype), + "mismatched_rows": mismatched, + "passed": same and not mismatched, + "problems": [] if same and not mismatched else [f"rows {mismatched[:8]} differ"], + } + + +def compare_routes(actual: dict, expected: dict, label: str) -> list: + """Every step of a route against another route of the same input.""" + results = [ + compare_exact("groups", actual["groups"], expected["groups"]), + compare_exact("ids", actual["ids"], expected["ids"]), + ] + for name, atol in RULES.items(): + if name == "weights": + continue + results.append(compare(name, actual[name], expected[name], atol)) + results.append( + compare( + "masked_selection_scores", + actual["masked_selection_scores"], + expected["masked_selection_scores"], + RULES["selection_scores"], + ) + ) + # Weights belong to the experts they weigh: compared on rows routed identically. + same_ids = (actual["ids"] == expected["ids"]).all(dim=-1) + results.append( + compare( + "weights", actual["weights"][same_ids], expected["weights"][same_ids], RULES["weights"] + ) + ) + for r in results: + r["against"] = label + return results + + +def assert_all_pass(results: list) -> None: + failed = [ + f"{r['against']}:{r['tensor']}: {'; '.join(r['problems'])}" + for r in results + if not r["passed"] + ] + assert not failed, "\n".join(failed) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/moe/test_modeling_v2_beam_grouped_topk_router.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/moe/test_modeling_v2_beam_grouped_topk_router.py new file mode 100644 index 000000000000..69a7599e3c67 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/moe/test_modeling_v2_beam_grouped_topk_router.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_grouped_topk_router entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_router_bf16__sm_103__tp1.ops.moe.beam_grouped_topk_router import ( # noqa: E501 + beam_grouped_topk_router as op, +) + +__extra_import_path__ = ["..", "../../../.."] +import _router_cases as cases # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_grouped_topk_router requires a CUDA device" + + +def _build(spec, seed): + weight, bias, hidden = cases.matrix_case(seed, spec["tokens"]) + kwargs = {k: spec[k] for k in ("n_group", "topk_group", "top_k", "eps")} + return (hidden.cuda(), weight.cuda(), bias.cuda()), kwargs + + +@pytest.mark.parametrize("seed", cases.SEEDS) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:24] for c in op.CELLS]) +def test_certified_cells(cell, seed, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Every step within the rules in `_router_cases.RULES`; groups and ids exactly, in order.""" + certifying(op) + args, kwargs = _build(cell.spec, seed) + out = op(*args, **kwargs) + ref = op.reference(*args, **kwargs) + assert set(out) == set(ref) + assert out["ids"].shape == (cell.spec["tokens"], cell.spec["top_k"]) + cases.assert_all_pass(cases.compare_routes(out, ref, "reference")) + + +@pytest.mark.parametrize( + "case,match", + [ + ("batched_3d", "hidden_states must be"), + ("bias_length_one", "correction_bias must be"), + ("fp64_bias", "correction_bias must be torch.float32"), + ("top_k_beyond_kept_groups", "cannot be chosen"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + (hidden, weight, bias), kwargs = _build(op.CELLS[0].spec, seed=17) + if case == "batched_3d": + hidden = hidden.unsqueeze(0) + elif case == "bias_length_one": + bias = bias[:1] + elif case == "fp64_bias": + bias = bias.double() + else: + kwargs["top_k"] = kwargs["topk_group"] * 32 + 1 + with pytest.raises(AssertionError, match=match): + with validating(op): + op(hidden, weight, bias, **kwargs) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/test_router.py b/tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/test_router.py new file mode 100644 index 000000000000..be5e7f2c7fdb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/test_router.py @@ -0,0 +1,912 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the Beam MoE router component of block 3: bf16 / sm_103 / tp1. + +Configuration under test, fixed in ``_router_cases.py`` and nowhere else: 256 experts in +8 groups of 32, 4 groups kept, 8 experts per token, hidden 5120; FP32 router weight and +correction bias, BF16 already-normalized input; one sm_103 GPU; eager; generated +parameters; seeds {17, 41} x tokens {1, 17, 129}, and the controlled fixtures there. + +What is compared. Every call goes through the component's ``forward``; the op's mapping +of that same call is captured, and each of its steps is compared with the op's own +``reference`` on the same input under the rules in ``_router_cases.RULES``. Groups and +expert ids are compared exactly, in order. Each fixture also states the decision it was +built to force, independently of either implementation. + +Invocation (fp32 needs the TF32 overrides removed before torch is imported):: + + 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 + +Optional environment: + +* ``BEAM_MOE_ROUTER_SOURCE_CAPTURES``: a directory of routes captured from the source + model's router: ``seed/params.pt`` (``weight``, ``bias``), ``seed/T.pt`` + (``hidden`` and ``native``, a dict of ``logits``, ``scores``, ``corrected``, + ``group_score``, ``top_groups``, ``masked``, ``ids``, ``weights``) and + ``fixtures/.pt`` (``hidden``, ``bias`` and the same names). The generated inputs + must then equal the captured ones bit for bit, the op's ``reference`` must reproduce + every captured tensor bit for bit, and the component is compared with the captures + under the same rules as with the reference. +* ``BEAM_MOE_ROUTER_EVIDENCE``: a directory that receives one JSON record per test case: + every comparison with its rule, hashes of the routed tensors, the torch calls and CUDA + kernels of the forward, and the device and software it ran on. + +Dependencies: a TensorRT-LLM build for sm_103. No model files. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import json +import os +import platform +import sys +import types +from pathlib import Path + +import pytest +import torch +from torch.overrides import TorchFunctionMode, resolve_name + +import tensorrt_llm +import tensorrt_llm._torch.models.modeling_beam as builtin_beam +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_router_bf16__sm_103__tp1.ops.moe.beam_grouped_topk_router import ( # noqa: E501 + BeamGroupedTopkRouter, + beam_grouped_topk_router, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_router_bf16__sm_103__tp1.router import ( # noqa: E501 + NORM_EPS, + BeamMoeRouterBf16Sm103Tp1, +) + +__extra_import_path__ = [".", "../../.."] +import _router_cases as cases # noqa: E402 +from _validating import validating # noqa: E402 + +assert torch.cuda.is_available(), "the Beam MoE router tests require a CUDA device" + +#: The router as the checkpoint configures it. Only these fields are read. +ROUTER_CONFIG = types.MappingProxyType( + dict( + hidden_size=5120, + num_experts=256, + num_experts_per_tok=8, + n_group=8, + topk_group=4, + norm_topk_prob=True, + mlp_layer_types=("dense", "dense", "dense", "sparse"), + ) +) +LAYER_IDX = 3 +SM = (10, 3) +#: Eager, one GPU, one call per case. +EXECUTION = types.MappingProxyType(dict(eager=True, cuda_graph=False, devices=1)) +OP_KWARGS = types.MappingProxyType( + dict(n_group=cases.GROUPS, topk_group=cases.TOPK_GROUP, top_k=cases.TOP_K, eps=cases.NORM_EPS) +) +#: Op names -> the names a source capture stores them under. +CAPTURE_NAMES = types.MappingProxyType( + dict( + logits="logits", + scores="scores", + selection_scores="corrected", + group_scores="group_score", + groups="top_groups", + masked_selection_scores="masked", + ids="ids", + weights="weights", + ) +) +#: Calls that move token data to the host, and with it a routing decision. +HOST_TRANSFERS = frozenset( + "torch.Tensor." + name + for name in ("cpu", "item", "tolist", "numpy", "__bool__", "__int__", "__float__", "__index__") +) + + +def router_config(**overrides) -> types.SimpleNamespace: + fields = dict(ROUTER_CONFIG) + fields["mlp_layer_types"] = list(fields["mlp_layer_types"]) + fields.update(overrides) + return types.SimpleNamespace(**fields) + + +@pytest.fixture(scope="module", autouse=True) +def no_tf32_override(): + """The router's fp32 steps are IEEE fp32; these variables substitute TF32 regardless of + what torch is told.""" + for var in ("TORCH_ALLOW_TF32_CUBLAS_OVERRIDE", "NVIDIA_TF32_OVERRIDE"): + assert var not in os.environ, ( + f"{var} is set and substitutes TF32 for fp32 GEMMs; run with `env -u {var}`" + ) + + +# ----------------------------------------------------------------------------- evidence + + +def _environment() -> dict: + device = torch.cuda.current_device() + return { + "device_name": torch.cuda.get_device_name(device), + "device_capability": list(torch.cuda.get_device_capability(device)), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "tensorrt_llm": tensorrt_llm.__version__, + "tensorrt_llm_file": tensorrt_llm.__file__, + "python": sys.version.split()[0], + "machine": platform.machine(), + "argv": sys.argv, + "pid": os.getpid(), + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "slurm_step_id": os.environ.get("SLURM_STEP_ID"), + "execution": dict(EXECUTION), + "ambient_math": { + "float32_matmul_precision": torch.get_float32_matmul_precision(), + "cuda_matmul_fp32_precision": torch.backends.cuda.matmul.fp32_precision, + "cuda_matmul_allow_tf32": bool(torch.backends.cuda.matmul.allow_tf32), + }, + "source_captures": os.environ.get("BEAM_MOE_ROUTER_SOURCE_CAPTURES"), + } + + +def _write_evidence(name: str, record: dict) -> None: + directory = os.environ.get("BEAM_MOE_ROUTER_EVIDENCE") + if not directory: + return + path = Path(directory) + path.mkdir(parents=True, exist_ok=True) + (path / f"{name}.json").write_text(json.dumps(record, indent=1, default=str) + "\n") + + +def tensor_sha256(t: torch.Tensor) -> str: + flat = t.detach().contiguous().cpu().reshape(-1) + h = hashlib.sha256() + h.update(f"{t.dtype}|{tuple(t.shape)}|".encode()) + h.update(flat.view(torch.uint8).numpy().tobytes()) + return h.hexdigest() + + +def route_hashes(route: dict) -> dict: + return {name: tensor_sha256(t) for name, t in sorted(route.items())} + + +# ----------------------------------------------------------------------------- source captures + + +def _captures_dir() -> Path | None: + directory = os.environ.get("BEAM_MOE_ROUTER_SOURCE_CAPTURES") + return Path(directory) if directory else None + + +def _as_route(captured: dict) -> dict: + return {name: captured[key] for name, key in CAPTURE_NAMES.items()} + + +def matrix_captures(seed: int, tokens: int) -> dict | None: + root = _captures_dir() + if root is None: + return None + params = torch.load(root / f"seed{seed}" / "params.pt", weights_only=True) + blob = torch.load(root / f"seed{seed}" / f"T{tokens}.pt", weights_only=True) + return dict( + weight=params["weight"], + bias=params["bias"], + hidden=blob["hidden"], + route=_as_route(blob["native"]), + ) + + +def fixture_captures(name: str) -> dict | None: + root = _captures_dir() + if root is None: + return None + blob = torch.load(root / "fixtures" / f"{name}.pt", weights_only=True) + return dict(bias=blob["bias"], hidden=blob["hidden"], route=_as_route(blob)) + + +def check_against_captures( + captured: dict | None, inputs: dict, route: dict, reference: dict +) -> list: + """With captures: same input bytes, the reference reproduces them bitwise, and the + component meets the rules against them. Without: nothing.""" + if captured is None: + return [] + results = [] + for name, tensor in inputs.items(): + same = torch.equal(tensor.cpu(), captured[name].cpu()) + results.append( + { + "tensor": f"input:{name}", + "against": "source_captures", + "rule": "bitwise", + "passed": same, + "problems": [] if same else [f"generated {name} differs from the capture"], + } + ) + for name, tensor in captured["route"].items(): + same = torch.equal(reference[name].cpu(), tensor.cpu()) + results.append( + { + "tensor": f"reference:{name}", + "against": "source_captures", + "rule": "bitwise", + "passed": same, + "problems": [] if same else [f"reference {name} is not the captured {name}"], + } + ) + device_route = {k: v.to(route["ids"].device) for k, v in captured["route"].items()} + results.extend(cases.compare_routes(route, device_route, "source_captures")) + return results + + +# ----------------------------------------------------------------------------- observing a call + + +@contextlib.contextmanager +def recording_op_calls(): + """Capture the mapping each op call returns, without changing what it returns.""" + calls: list[dict] = [] + original = BeamGroupedTopkRouter.__call__ + + def recorded(self, *args, **kwargs): + out = original(self, *args, **kwargs) + calls.append(out) + return out + + BeamGroupedTopkRouter.__call__ = recorded + try: + yield calls + finally: + BeamGroupedTopkRouter.__call__ = original + + +@contextlib.contextmanager +def builtin_router_forbidden(): + """The built-in Beam router must not compute any part of this component's routing.""" + saved = {name: getattr(builtin_beam.BeamRouter, name) for name in ("forward", "route")} + + def refuse(*args, **kwargs): + raise AssertionError("the built-in BeamRouter was called") + + for name in saved: + setattr(builtin_beam.BeamRouter, name, refuse) + try: + yield + finally: + for name, fn in saved.items(): + setattr(builtin_beam.BeamRouter, name, fn) + + +def _tensors_of(value): + if isinstance(value, torch.Tensor): + return [value] + if isinstance(value, (tuple, list)): + return [t for v in value for t in _tensors_of(v)] + if isinstance(value, dict): + return [t for v in value.values() for t in _tensors_of(v)] + return [] + + +class TorchCallTrace(TorchFunctionMode): + """Every torch call made inside the block, with the device and dtype of what it returned, + and the fp32 matmul and autocast state at each projection.""" + + def __init__(self): + super().__init__() + self.calls: list[dict] = [] + + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + name = resolve_name(func) or getattr(func, "__qualname__", None) or repr(func) + entry = {"func": name} + if name == "torch.nn.functional.linear": + entry["math"] = { + "cuda_matmul_fp32_precision": torch.backends.cuda.matmul.fp32_precision, + "cuda_matmul_allow_tf32": bool(torch.backends.cuda.matmul.allow_tf32), + "autocast_cuda": torch.is_autocast_enabled("cuda"), + } + entry["operand_dtypes"] = [str(t.dtype) for t in _tensors_of(args)] + out = func(*args, **kwargs) + produced = _tensors_of(out) + if produced: + entry["outputs"] = [[str(t.device), str(t.dtype), list(t.shape)] for t in produced] + self.calls.append(entry) + return out + + +def loaded_router(weight: torch.Tensor, bias: torch.Tensor) -> BeamMoeRouterBf16Sm103Tp1: + router = BeamMoeRouterBf16Sm103Tp1(router_config(), LAYER_IDX, device="cuda") + router.load_weights(weight.cuda(), bias.cuda()) + return router + + +def routed(router: BeamMoeRouterBf16Sm103Tp1, hidden: torch.Tensor) -> tuple: + """(weights, ids, the op's mapping of that same forward call).""" + with builtin_router_forbidden(), validating(beam_grouped_topk_router): + with recording_op_calls() as calls: + weights, ids = router(hidden) + assert len(calls) == 1, f"forward made {len(calls)} op calls, expected 1" + (route,) = calls + assert weights is route["weights"] and ids is route["ids"], ( + "forward did not return the op call's own weights and ids" + ) + return weights, ids, route + + +def reference_route(router: BeamMoeRouterBf16Sm103Tp1, hidden: torch.Tensor) -> dict: + return beam_grouped_topk_router.reference( + hidden, router.weight, router.e_score_correction_bias, **OP_KWARGS + ) + + +def check_output_contract(weights, ids, hidden) -> list: + tokens = hidden.shape[0] + problems = [] + if tuple(weights.shape) != (tokens, 8) or weights.dtype != torch.float32: + problems.append(f"weights {tuple(weights.shape)} {weights.dtype}") + if tuple(ids.shape) != (tokens, 8) or ids.dtype != torch.int64: + problems.append(f"ids {tuple(ids.shape)} {ids.dtype}") + if weights.device != hidden.device or ids.device != hidden.device: + problems.append(f"outputs on {weights.device}/{ids.device}, input on {hidden.device}") + return [ + { + "tensor": "output_contract", + "against": "contract", + "rule": "[T, 8] float32 weights and int64 ids on the input device", + "passed": not problems, + "problems": problems, + } + ] + + +def check_invariants(route: dict, sum_to_one: bool) -> list: + """Per row: eight distinct ids in range, each in a selected group; finite, nonnegative + weights; with `sum_to_one`, |sum(weights) - 1| <= NORM_ATOL.""" + ids, groups, weights = route["ids"], route["groups"], route["weights"] + distinct = (ids.sort(dim=-1).values.diff(dim=-1) != 0).all(dim=-1) + in_range = ((ids >= 0) & (ids < cases.EXPERTS)).all(dim=-1) + member = ((ids // cases.PER_GROUP).unsqueeze(-1) == groups.unsqueeze(1)).any(dim=-1).all(-1) + excluded = route["masked_selection_scores"] == float("-inf") + expert_selected_group = ( + (torch.arange(cases.EXPERTS, device=ids.device) // cases.PER_GROUP).view(1, -1, 1) + == groups.unsqueeze(1) + ).any(dim=-1) + mask_exact = (excluded == ~expert_selected_group).all(dim=-1) + finite = (torch.isfinite(weights) & (weights >= 0)).all(dim=-1) + deviation = (weights.sum(dim=-1).double() - 1.0).abs() + checks = { + "eight_distinct_ids": distinct, + "ids_in_range": in_range, + "ids_in_selected_groups": member, + "excluded_exactly_outside_selected_groups": mask_exact, + "weights_finite_nonnegative": finite, + } + if sum_to_one: + checks["weights_sum_to_one"] = deviation <= cases.NORM_ATOL + results = [] + for name, ok in checks.items(): + bad = [int(i) for i in torch.nonzero(~ok).flatten().cpu()] + results.append( + { + "tensor": name, + "against": "invariant", + "rule": f"|sum - 1| <= {cases.NORM_ATOL:g}" + if name == "weights_sum_to_one" + else name, + "passed": not bad, + "problems": [f"rows {bad[:8]}"] if bad else [], + } + ) + results[-1]["max_abs_sum_minus_one"] = float(deviation.max()) + return results + + +# ----------------------------------------------------------------------------- construction and loading + + +def test_construction_allocates_only_the_router_state() -> None: + """Construction allocates exactly the two FP32 router tensors: no expert, attention, + cache or model state, and no submodule.""" + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + router = BeamMoeRouterBf16Sm103Tp1(router_config(), LAYER_IDX, device="cuda") + torch.cuda.synchronize() + allocated = torch.cuda.memory_allocated() - before + params = {n: (tuple(p.shape), p.dtype, p.requires_grad) for n, p in router.named_parameters()} + buffers = {n: (tuple(b.shape), b.dtype) for n, b in router.named_buffers()} + elements = sum(p.numel() for p in router.parameters()) + sum( + b.numel() for b in router.buffers() + ) + record = { + "test": "test_construction_allocates_only_the_router_state", + "parameters": params, + "buffers": buffers, + "submodules": [n for n, _ in router.named_modules() if n], + "elements": elements, + "cuda_bytes_allocated": allocated, + "environment": _environment(), + } + _write_evidence("construction", record) + assert params == {"weight": ((256, 5120), torch.float32, False)} + assert buffers == {"e_score_correction_bias": ((256,), torch.float32)} + assert record["submodules"] == [] + assert elements == 256 * 5120 + 256 + assert allocated == (256 * 5120 + 256) * 4, allocated + assert router.weight.device.type == "cuda" + assert torch.cuda.get_device_capability(router.weight.device) == SM + + +@pytest.mark.parametrize( + "overrides,layer_idx,match", + [ + (dict(n_group=4), LAYER_IDX, "unsupported n_group"), + (dict(topk_group=3), LAYER_IDX, "unsupported topk_group"), + (dict(num_experts_per_tok=6), LAYER_IDX, "unsupported num_experts_per_tok"), + (dict(num_experts=128), LAYER_IDX, "unsupported num_experts"), + (dict(hidden_size=4096), LAYER_IDX, "unsupported hidden_size"), + (dict(norm_topk_prob=False), LAYER_IDX, "unsupported norm_topk_prob"), + (dict(mlp_layer_types=["dense"] * 4), LAYER_IDX, "sparse"), + ({}, 0, "only the routers of sparse blocks"), + ], +) +def test_construction_rejects_unsupported_configs(overrides, layer_idx, match) -> None: + with pytest.raises(ValueError, match=match): + BeamMoeRouterBf16Sm103Tp1(router_config(**overrides), layer_idx, device="cuda") + + +def test_construction_rejects_a_cpu_device() -> None: + with pytest.raises(ValueError, match="sm_103 CUDA devices"): + BeamMoeRouterBf16Sm103Tp1(router_config(), LAYER_IDX, device="cpu") + + +@pytest.mark.parametrize("seed", cases.SEEDS) +def test_loading_copies_both_tensors_exactly(seed) -> None: + weight, bias, _ = cases.matrix_case(seed, 1) + router = BeamMoeRouterBf16Sm103Tp1(router_config(), LAYER_IDX, device="cuda") + router.load_weights(weight, bias) + _write_evidence( + f"loading-seed{seed}", + { + "test": "test_loading_copies_both_tensors_exactly", + "weight": [ + list(router.weight.shape), + str(router.weight.dtype), + tensor_sha256(router.weight), + ], + "bias": [ + list(router.e_score_correction_bias.shape), + str(router.e_score_correction_bias.dtype), + tensor_sha256(router.e_score_correction_bias), + ], + "generated": [tensor_sha256(weight), tensor_sha256(bias)], + }, + ) + assert router.weight.dtype == torch.float32 and tuple(router.weight.shape) == (256, 5120) + assert router.e_score_correction_bias.dtype == torch.float32 + assert tuple(router.e_score_correction_bias.shape) == (256,) + assert torch.equal(router.weight.cpu(), weight) + assert torch.equal(router.e_score_correction_bias.cpu(), bias) + + +@pytest.mark.parametrize( + "case,error,match", + [ + ("weight_rows", ValueError, r"weight must be \[256, 5120\]"), + ("weight_width", ValueError, r"weight must be \[256, 5120\]"), + ("weight_transposed", ValueError, r"weight must be \[256, 5120\]"), + ("bias_length", ValueError, r"correction_bias must be \[256\]"), + ("bias_length_one", ValueError, r"correction_bias must be \[256\]"), + ("weight_bf16", TypeError, "weight must be torch.float32"), + ("bias_fp64", TypeError, "correction_bias must be torch.float32"), + ("weight_not_tensor", TypeError, "weight must be a torch.Tensor"), + ], +) +def test_loading_rejects_invalid_parameters(case, error, match) -> None: + """A rejected load raises before copying either tensor: the router keeps its state.""" + weight, bias, _ = cases.matrix_case(17, 1) + router = loaded_router(weight, bias) + good_weight, good_bias = router.weight.clone(), router.e_score_correction_bias.clone() + other_weight, other_bias, _ = cases.matrix_case(41, 1) + bad = { + "weight_rows": (other_weight[:255], other_bias), + "weight_width": (other_weight[:, :5119], other_bias), + "weight_transposed": (other_weight.t().contiguous(), other_bias), + "bias_length": (other_weight, other_bias[:255]), + "bias_length_one": (other_weight, other_bias[:1]), + "weight_bf16": (other_weight.to(torch.bfloat16), other_bias), + "bias_fp64": (other_weight, other_bias.double()), + "weight_not_tensor": ([0.0] * 5120, other_bias), + }[case] + with pytest.raises(error, match=match): + router.load_weights(*bad) + assert torch.equal(router.weight, good_weight) + assert torch.equal(router.e_score_correction_bias, good_bias) + + +@pytest.mark.parametrize( + "case,error,match", + [ + ("width_5119", ValueError, r"hidden_states must be \[tokens, 5120\]"), + ("width_5121", ValueError, r"hidden_states must be \[tokens, 5120\]"), + ("batched_3d", ValueError, r"hidden_states must be \[tokens, 5120\]"), + ("one_dim", ValueError, r"hidden_states must be \[tokens, 5120\]"), + ("fp32_input", TypeError, "hidden_states must be torch.bfloat16"), + ("fp16_input", TypeError, "hidden_states must be torch.bfloat16"), + ("no_tokens", ValueError, "at least one token"), + ("cpu_input", ValueError, "router's device"), + ("not_a_tensor", TypeError, "hidden_states must be a torch.Tensor"), + ("not_loaded", RuntimeError, "load_weights must be called"), + ], +) +def test_forward_rejects_invalid_input(case, error, match) -> None: + weight, bias, hidden = cases.matrix_case(17, 17) + if case == "not_loaded": + router = BeamMoeRouterBf16Sm103Tp1(router_config(), LAYER_IDX, device="cuda") + else: + router = loaded_router(weight, bias) + hidden = hidden.cuda() + bad = { + "width_5119": hidden[:, :5119], + "width_5121": torch.cat([hidden, hidden[:, :1]], dim=1), + "batched_3d": hidden.unsqueeze(0), + "one_dim": hidden[0], + "fp32_input": hidden.float(), + "fp16_input": hidden.half(), + "no_tokens": hidden[:0], + "cpu_input": hidden.cpu(), + "not_a_tensor": [0.0] * 5120, + "not_loaded": hidden, + }[case] + with pytest.raises(error, match=match): + router(bad) + + +# ----------------------------------------------------------------------------- precision + + +def test_precision_path_is_fp32_without_tf32_or_autocast() -> None: + """Projection, sigmoid, selection scores, group scores and normalization are fp32, the + projection is IEEE fp32 even when the caller has TF32 or bf16 autocast on, and the + caller's settings are restored afterwards.""" + weight, bias, hidden = cases.matrix_case(17, 17) + router = loaded_router(weight, bias) + hidden = hidden.cuda() + plain_weights, plain_ids, plain = routed(router, hidden) + + saved = (torch.backends.cuda.matmul.allow_tf32, torch.backends.cuda.matmul.fp32_precision) + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cuda.matmul.fp32_precision = "tf32" + try: + ambient = (torch.backends.cuda.matmul.allow_tf32, torch.backends.cuda.matmul.fp32_precision) + with torch.autocast("cuda", dtype=torch.bfloat16), TorchCallTrace() as trace: + weights, ids, route = routed(router, hidden) + after = (torch.backends.cuda.matmul.allow_tf32, torch.backends.cuda.matmul.fp32_precision) + finally: + torch.backends.cuda.matmul.allow_tf32, torch.backends.cuda.matmul.fp32_precision = saved + + # TF32 operands are ~1e-3 off; an IEEE fp32 product of these operands is ~1e-6 off fp64. + fp64_logits = hidden.double() @ router.weight.double().t() + fp64_error = float((route["logits"].double() - fp64_logits).abs().max()) + linear = [c for c in trace.calls if c["func"] == "torch.nn.functional.linear"] + dtypes = {name: str(t.dtype) for name, t in route.items()} + _write_evidence( + "precision", + { + "test": "test_precision_path_is_fp32_without_tf32_or_autocast", + "ambient_during_call": {"tf32": list(ambient), "autocast_bf16": True}, + "ambient_after_call": list(after), + "projection_calls": linear, + "route_dtypes": dtypes, + "logits_vs_fp64_max_abs": fp64_error, + "bitwise_equal_to_plain_call": {k: torch.equal(route[k], plain[k]) for k in route}, + "torch_calls": [c["func"] for c in trace.calls], + "environment": _environment(), + }, + ) + assert ambient == (True, "tf32") and after == (True, "tf32"), (ambient, after) + assert len(linear) == 1, linear + assert linear[0]["math"] == { + "cuda_matmul_fp32_precision": "ieee", + "cuda_matmul_allow_tf32": False, + "autocast_cuda": False, + }, linear[0] + assert linear[0]["operand_dtypes"] == ["torch.float32", "torch.float32"], linear[0] + assert linear[0]["outputs"][0][1] == "torch.float32" + for name in ("logits", "scores", "selection_scores", "group_scores", "weights"): + assert dtypes[name] == "torch.float32", (name, dtypes[name]) + assert dtypes["masked_selection_scores"] == "torch.float32" + assert dtypes["ids"] == dtypes["groups"] == "torch.int64" + assert fp64_error <= 1e-5, f"projection is not IEEE fp32: {fp64_error:.3e} from fp64" + for name in route: + assert torch.equal(route[name], plain[name]), f"{name} changed under TF32/autocast" + assert torch.equal(weights, plain_weights) and torch.equal(ids, plain_ids) + + +# ----------------------------------------------------------------------------- the fixed matrix + + +@pytest.mark.parametrize("seed,tokens", cases.CASES, ids=cases.CASE_IDS) +def test_matrix(seed, tokens) -> None: + """One case of the fixed matrix: the output contract, every step against the reference + (and the source captures when given), the routing invariants, and the evidence that + the routing ran on the GPU in this component's op with no host synchronization.""" + weight, bias, hidden_cpu = cases.matrix_case(seed, tokens) + router = loaded_router(weight, bias) + hidden = hidden_cpu.cuda() + torch.cuda.synchronize() + + activities = [torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] + with torch.profiler.profile(activities=activities) as prof: + with TorchCallTrace() as trace: + previous = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("error") + try: + weights, ids, route = routed(router, hidden) + finally: + torch.cuda.set_sync_debug_mode(previous) + torch.cuda.synchronize() + kernels = [e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA] + + reference = reference_route(router, hidden) + results = check_output_contract(weights, ids, hidden) + results += cases.compare_routes(route, reference, "reference") + results += check_invariants(route, sum_to_one=True) + results += check_against_captures( + matrix_captures(seed, tokens), + {"weight": weight, "bias": bias, "hidden": hidden_cpu}, + route, + reference, + ) + + host = [c["func"] for c in trace.calls if c["func"] in HOST_TRANSFERS] + off_device = [ + c for c in trace.calls for o in c.get("outputs", []) if not o[0].startswith("cuda") + ] + called = {c["func"] for c in trace.calls} + routing_calls = { + "torch.nn.functional.linear", + "torch.Tensor.sigmoid", + "torch.Tensor.topk", + "torch.Tensor.scatter_", + "torch.Tensor.masked_fill", + "torch.Tensor.gather", + "torch.Tensor.sum", + } + _write_evidence( + f"matrix-seed{seed}-T{tokens}", + { + "test": "test_matrix", + "case": {"seed": seed, "tokens": tokens}, + "component": f"{type(router).__module__}.{type(router).__qualname__}", + "op": f"{BeamGroupedTopkRouter.__module__}.{BeamGroupedTopkRouter.__qualname__}", + "rules": dict(cases.RULES, norm_atol=cases.NORM_ATOL), + "results": results, + "hashes": route_hashes(route), + "inputs": { + "weight": tensor_sha256(weight), + "bias": tensor_sha256(bias), + "hidden": tensor_sha256(hidden_cpu), + }, + "ids": route["ids"].tolist(), + "groups": route["groups"].tolist(), + "torch_calls": trace.calls, + "host_transfers": host, + "off_device_outputs": off_device, + "sync_debug_mode": "error", + "cuda_kernels": kernels, + "environment": _environment(), + }, + ) + assert type(router).__module__.endswith("moe_router_bf16__sm_103__tp1.router") + assert torch.cuda.get_device_capability(hidden.device) == SM + assert not host, f"routing moved token data to the host: {host}" + assert not off_device, f"routing produced tensors off the GPU: {off_device[:3]}" + assert routing_calls <= called, sorted(routing_calls - called) + assert kernels, "no CUDA kernel ran" + cases.assert_all_pass(results) + + +# ----------------------------------------------------------------------------- controlled fixtures + + +def route_fixture(name: str) -> tuple[dict, dict, list]: + """(component route, reference route, comparisons) for one fixture, captures included.""" + spec = cases.fixtures()[name] + weight, bias, hidden_cpu = cases.fixture_tensors(spec) + router = loaded_router(weight, bias) + hidden = hidden_cpu.cuda() + _, _, route = routed(router, hidden) + reference = reference_route(router, hidden) + results = cases.compare_routes(route, reference, "reference") + results += check_invariants(route, sum_to_one=False) + results += check_against_captures( + fixture_captures(name), {"bias": bias, "hidden": hidden_cpu}, route, reference + ) + exact = torch.equal(route["logits"].cpu(), spec["logits"].float()) + results.append( + { + "tensor": "fixture_logits", + "against": "fixture", + "rule": "the one-hot rows reproduce the fixture's logits exactly", + "passed": exact, + "problems": [] if exact else ["logits differ from the fixture's"], + } + ) + return route, reference, results + + +def _record_fixture(name: str, test: str, route: dict, results: list, extra: dict) -> None: + _write_evidence( + f"fixture-{name}", + { + "test": test, + "fixture": name, + "intent": cases.fixtures()[name]["intent"], + "rows": list(cases.fixtures()[name]["rows"]), + "ids": route["ids"].tolist(), + "groups": route["groups"].tolist(), + "weights": route["weights"].tolist(), + "row_sums": route["weights"].sum(dim=-1).tolist(), + "hashes": route_hashes(route), + "results": results, + **extra, + "environment": _environment(), + }, + ) + + +def test_bias_selects_but_never_weights() -> None: + """The correction bias changes which experts are chosen, and the returned weights are + still the normalized unbiased scores of the chosen experts.""" + zero, _, zero_results = route_fixture("bias_zero") + shift, _, shift_results = route_fixture("bias_shift") + ids, scores = shift["ids"], shift["scores"] + unbiased = scores.gather(-1, ids) + unbiased = unbiased / (unbiased.sum(dim=-1, keepdim=True) + NORM_EPS) + biased = shift["selection_scores"].gather(-1, ids) + biased = biased / (biased.sum(dim=-1, keepdim=True) + NORM_EPS) + selected_bias = cases.fixtures()["bias_shift"]["bias"][ids[0].cpu()] + leak_gap = float((biased - shift["weights"]).abs().max()) + extra = { + "ids_without_bias": zero["ids"].tolist(), + "ids_with_bias": ids.tolist(), + "groups_with_bias": shift["groups"].tolist(), + "selected_bias": selected_bias.tolist(), + "weights_vs_normalized_unbiased": cases.compare( + "weights", shift["weights"], unbiased, cases.RULES["weights"] + ), + "weights_vs_normalized_biased_max_abs": leak_gap, + } + _record_fixture("bias_zero", "test_bias_selects_but_never_weights", zero, zero_results, {}) + _record_fixture( + "bias_shift", "test_bias_selects_but_never_weights", shift, shift_results, extra + ) + cases.assert_all_pass(zero_results + shift_results) + assert zero["groups"].tolist() == [[0, 1, 2, 3]] + assert zero["ids"].tolist() == [[0, 1, 32, 33, 64, 65, 96, 97]] + assert shift["groups"].tolist() == [[6, 0, 1, 2]] + assert ids.tolist() == [[193, 192, 0, 1, 32, 33, 64, 65]] + assert torch.equal(zero["logits"], shift["logits"]) and torch.equal(zero["scores"], scores) + assert bool((selected_bias != 0).any()) and len(set(selected_bias.tolist())) > 1 + assert extra["weights_vs_normalized_unbiased"]["passed"], extra[ + "weights_vs_normalized_unbiased" + ] + assert leak_gap > cases.RULES["weights"], f"bias in the weights would be invisible: {leak_gap}" + + +def test_group_boundary_moves_with_the_top_two_sum() -> None: + """Only the challenger group's second expert moves, and the group crosses the 4th/5th + boundary; the group holding the single best expert stays out.""" + route, _, results = route_fixture("group_boundary") + group_scores = route["group_scores"] + _record_fixture( + "group_boundary", + "test_group_boundary_moves_with_the_top_two_sum", + route, + results, + {"group_scores": group_scores.tolist()}, + ) + cases.assert_all_pass(results) + assert route["groups"].tolist() == [[0, 1, 2, 3], [0, 1, 2, 4]] + assert route["ids"].tolist() == [ + [0, 1, 32, 33, 64, 65, 96, 97], + [0, 1, 32, 33, 64, 65, 129, 128], + ] + # Group 4 scores 0.66 + 0.62 then 0.66 + 0.67 against group 3's 0.655 + 0.645. + assert float(group_scores[0, 4]) < float(group_scores[0, 3]) < float(group_scores[1, 4]) + assert 5 not in route["groups"].flatten().tolist() + + +def test_exclusion_masks_every_expert_of_dropped_groups() -> None: + """The best single expert sits in a group that loses; it and every other expert of the + four dropped groups are -inf, although every retained selection score is negative.""" + route, _, results = route_fixture("exclusion") + ids = route["ids"] + selected_scores = route["selection_scores"].gather(-1, ids) + spike_score = float(route["selection_scores"][0, cases.EXCLUSION_SPIKE]) + _record_fixture( + "exclusion", + "test_exclusion_masks_every_expert_of_dropped_groups", + route, + results, + {"spike": cases.EXCLUSION_SPIKE, "spike_selection_score": spike_score}, + ) + cases.assert_all_pass(results) + assert sorted(route["groups"][0].tolist()) == [0, 1, 2, 3] + assert cases.EXCLUSION_SPIKE not in ids[0].tolist() + assert spike_score == float(route["selection_scores"].max()) + assert bool((selected_scores < 0).all()) + assert all(e // cases.PER_GROUP in route["groups"][0].tolist() for e in ids[0].tolist()) + + +def test_expert_boundary_moves_one_expert() -> None: + """Expert 2 crosses the 8th/9th boundary within an unchanged group selection: it enters + last and expert 97 leaves, every other position unchanged.""" + route, _, results = route_fixture("expert_boundary") + _record_fixture("expert_boundary", "test_expert_boundary_moves_one_expert", route, results, {}) + cases.assert_all_pass(results) + assert route["groups"].tolist() == [[0, 1, 2, 3], [0, 1, 2, 3]] + assert route["ids"].tolist() == [ + [0, 1, 32, 33, 64, 65, 96, 97], + [0, 1, 32, 33, 64, 65, 96, 2], + ] + + +def test_ties_follow_the_pinned_topk_order() -> None: + """Exactly tied scores are routed in the order `torch.topk(..., sorted=True)` returns + -- the source's order, stated in `_router_cases.TIE_ROUTES` -- whether a row is routed + in a batch or alone. Ids are compared in order, never sorted.""" + route, _, results = route_fixture("ties") + spec = cases.fixtures()["ties"] + weight, bias, hidden = cases.fixture_tensors(spec) + router = loaded_router(weight, bias) + alone = {} + for i, name in enumerate(spec["rows"]): + _, _, single = routed(router, hidden[i : i + 1].cuda()) + alone[name] = (single["groups"][0].tolist(), single["ids"][0].tolist()) + batched = { + name: (route["groups"][i].tolist(), route["ids"][i].tolist()) + for i, name in enumerate(spec["rows"]) + } + _record_fixture( + "ties", + "test_ties_follow_the_pinned_topk_order", + route, + results, + {"batched": batched, "alone": alone, "pinned": cases.TIE_ROUTES}, + ) + cases.assert_all_pass(results) + for name, pinned in cases.TIE_ROUTES.items(): + assert batched[name] == (list(pinned[0]), list(pinned[1])), (name, batched[name]) + assert alone[name] == batched[name], (name, alone[name]) + + +def test_epsilon_is_the_source_denominator() -> None: + """Weights divide by sum + 1e-20: a row whose scores underflow to 0 gets zero weights, + not NaN, and a row whose selected scores sum to ~1e-20 sums to S / (S + 1e-20), well + below 1 -- a denominator without the term, or with 1e-12 or 1e-6, gives 1 or ~0.""" + route, _, results = route_fixture("epsilon") + selected = route["scores"].gather(-1, route["ids"]) + exact = torch.sigmoid(route["logits"].double()).gather(-1, route["ids"]).sum(dim=-1) + expected_sum = (exact / (exact + 1e-20)).tolist() + sums = route["weights"].sum(dim=-1).double().tolist() + _record_fixture( + "epsilon", + "test_epsilon_is_the_source_denominator", + route, + results, + { + "selected_score_sums": selected.sum(dim=-1).tolist(), + "expected_row_sums_fp64": expected_sum, + }, + ) + cases.assert_all_pass(results) + assert route["ids"].tolist() == [cases.EPSILON_IDS, cases.EPSILON_IDS] + assert route["groups"].tolist() == [[0, 1, 2, 3], [0, 1, 2, 3]] + assert bool((route["weights"][0] == 0).all()), route["weights"][0] + assert float(selected[1].sum()) > 0 + assert sums[1] < 0.9, sums[1] + assert abs(sums[1] - expected_sum[1]) <= cases.NORM_ATOL, (sums[1], expected_sum[1]) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_beam_logit_softcap.py b/tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_beam_logit_softcap.py new file mode 100644 index 000000000000..2f5bdaa6a948 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_beam_logit_softcap.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the six-block model's beam_logit_softcap entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.small_model_prefill_bf16__sm_103__tp1.ops.activation.beam_logit_softcap import ( # noqa: E501 + beam_logit_softcap as op, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_logit_softcap requires a CUDA device" + +#: Signed pre-softcap values from zero to far beyond saturation, -0 included. +SATURATION = (0.0, -0.0, 2.0**-10, -(2.0**-10), 0.5, -0.5, 1.0, -1.0, 5.0, -5.0, 10.0, -10.0) +SATURATION += (29.0, -29.0, 30.0, -30.0, 60.0, -60.0, 150.0, -150.0, 300.0, -300.0, 1024.0, -1024.0) +SATURATION += (32768.0, -32768.0, 3.0e38, -3.0e38) + + +def _build(spec, seed): + if spec.get("saturating"): + return torch.tensor(SATURATION, dtype=torch.float32, device="cuda").to(spec["dtype"])[None] + g = torch.Generator(device="cuda").manual_seed(seed) + logits = 2.0 * torch.randn(spec["tokens"], spec["vocab"], generator=g, device="cuda") + return logits.to(spec["dtype"]) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bitwise: the same three fp32 steps; finite, within the cap, fp32.""" + certifying(op) + logits = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(logits, cell.spec["cap"]) + assert out.shape == logits.shape and out.dtype == torch.float32 + torch.testing.assert_close(out, op.reference(logits, cell.spec["cap"]), rtol=0.0, atol=0.0) + assert bool(torch.isfinite(out).all()) and float(out.abs().max()) <= cell.spec["cap"] + + +def test_saturation_and_signs() -> None: + """Large magnitudes map to exactly +-30, signs are kept (-0 stays -0), and the result + tracks the float64 expression within a few fp32 units of the cap.""" + logits = _build(dict(saturating=True, dtype=torch.bfloat16), 0) + out = op(logits, 30.0)[0] + ref = torch.tanh(logits[0].double() / 30.0) * 30.0 + assert float((out.double() - ref).abs().max()) <= 8 * 30.0 * 2.0**-24 + big = logits[0].float().abs() >= 300.0 + assert torch.equal(out[big].abs(), torch.full_like(out[big], 30.0)) + assert torch.equal(torch.signbit(out), torch.signbit(logits[0].float())) + + +def test_once_is_not_twice_or_none() -> None: + """Applying the cap twice, or not at all, changes ordinary unsaturated logits.""" + logits = _build(dict(tokens=17, vocab=151936, dtype=torch.bfloat16), 5) + once = op(logits, 30.0) + assert not torch.equal(once, logits.float()) + assert not torch.equal(once, op(once.to(torch.bfloat16), 30.0)) + + +@pytest.mark.parametrize( + "case,match", [("int_cap", "positive float"), ("fp32_logits", "bf16 head output")] +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + logits = _build(dict(tokens=1, vocab=151936, dtype=torch.bfloat16), 9) + cap = 30.0 + if case == "int_cap": + cap = 30 + else: + logits = logits.float() + with pytest.raises(AssertionError, match=match): + with validating(op): + op(logits, cap) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_cached_generation.py b/tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_cached_generation.py new file mode 100644 index 000000000000..fddde53713ab --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_cached_generation.py @@ -0,0 +1,733 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for cached greedy generation and prefix reuse on the Beam six-block model. + +Configuration under test: the six-block model of ``test_small_model_prefill.py`` (bf16, sm_103, +eager TP1, generated weights) driven by ``generation.BeamSmallModelGenerator`` -- one +``KVCacheManagerV2`` with whole-page block reuse, 32-token pages and a 1024-token capacity, +TRTLLM attention through the catalog ``thop_attention`` op, and the existing ``TorchSampler``. + +What is checked, without model files or captured data: + +* the request lifecycle -- create, prefill, decode, sampled tokens and finish reason, + release -- with its counters, and that nothing runs after the finishing token; +* that a decode step appends exactly its own keys and values, one row per layer at the + token's absolute position, and changes no other written slot; +* prefix reuse through the manager: a matching prefix is claimed page by page and only the + rest of the prompt is computed, a different first token misses, and two live requests + sharing a prefix never write each other's slots; +* the sampler's greedy choice on identical logits (lowest index among exact maxima); +* the sliding window at the cache level: at query position 512 the sliding layer excludes + key 0 exactly and the full layer keeps it; +* rejections before any kernel runs: malformed or foreign request state, capacity overflow, + positions that do not continue the cache, a step that does not begin where the request's + cache history ends (rewound into cached tokens -- a prefix reused from another request + included -- or leaving a gap), a cached prefix without paged-context FMHA and a block-offset + table that does not address the request's pages, each leaving every written slot unchanged; + and weights that change -- in place or by replacing a parameter -- while a generator holds a + cache computed from them, refused before the step looks up a prefix or allocates. + +Numerical agreement with the source model is established separately. + +Invocation (fp32 steps need the TF32 overrides removed before torch is imported; the +generation-phase attention kernels compile on first use and need the installation's ``pip`` +on ``PATH``):: + + env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE python -m pytest \\ + tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1 + +Dependencies: a TensorRT-LLM build for sm_103 with flashinfer importable and about 130 GiB of +free device memory (generated state and loaded model). No model files. +""" + +from __future__ import annotations + +import contextlib +import gc +import math +import types +import zlib +from collections import OrderedDict + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + thop_attention, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam import _attention +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.small_model_prefill_bf16__sm_103__tp1 import ( # noqa: E501 + generation, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.small_model_prefill_bf16__sm_103__tp1 import ( # noqa: E501 + model as model_module, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.small_model_prefill_bf16__sm_103__tp1.generation import ( # noqa: E501 + BeamSmallModelGenerator, + Lifecycle, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState +from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.pyexecutor.seq_slot_manager import SeqSlotManager +from tensorrt_llm.bindings.executor import FinishReason + +assert torch.cuda.is_available(), "the Beam six-block model requires a CUDA device" + +pytestmark = pytest.mark.high_cuda_memory + +VOCAB = 151936 +TPB = 32 + +CONFIG = types.MappingProxyType( + dict( + num_hidden_layers=6, + hidden_size=5120, + intermediate_size=20480, + moe_intermediate_size=2560, + num_attention_heads=48, + num_key_value_heads=8, + head_dim=128, + vocab_size=VOCAB, + rms_norm_eps=1e-6, + hidden_act="silu", + attention_bias=False, + attention_dropout=0.0, + sliding_window=512, + layer_types=("sliding_attention",) * 5 + ("full_attention",), + mlp_layer_types=("dense",) * 3 + ("sparse",) * 3, + rope_parameters={ + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + }, + num_experts=256, + num_experts_per_tok=8, + n_group=8, + topk_group=4, + norm_topk_prob=True, + routed_scaling_factor=2.5, + tie_word_embeddings=False, + final_logit_softcapping=30.0, + max_position_embeddings=1048576, + ) +) + + +def model_config() -> types.SimpleNamespace: + fields = dict(CONFIG) + fields["layer_types"] = list(fields["layer_types"]) + fields["mlp_layer_types"] = list(fields["mlp_layer_types"]) + fields["rope_parameters"] = { + k: (dict(v) if v else v) for k, v in fields["rope_parameters"].items() + } + return types.SimpleNamespace(**fields) + + +def _half_width(name: str, shape: tuple) -> tuple[float, float]: + if name == model_module.EMBED: + return math.sqrt(3.0), 0.0 + if name.endswith("norm.weight") or name.endswith("layernorm.weight"): + return 0.4, 1.0 + if name.endswith("gate_bias"): + return 0.75 * math.sqrt(3.0), 0.0 + if name.endswith("e_score_correction_bias"): + return 0.05 * math.sqrt(3.0), 0.0 + return math.sqrt(3.0 / shape[-1]), 0.0 + + +def generate_state(seed: int) -> "OrderedDict[str, torch.Tensor]": + """All 2394 tensors, each from its own generator keyed by seed and name, fp32 rounded once.""" + state = OrderedDict() + for name, (shape, dtype) in model_module.weight_schema(model_config()).items(): + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(f"{seed}|{name}".encode())) + half, offset = _half_width(name, shape) + t = torch.rand(shape, generator=g, device="cuda", dtype=torch.float32) + state[name] = (t.mul_(2.0 * half).sub_(half).add_(offset)).to(dtype) + return state + + +def token_stream(name: str, length: int) -> list: + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(f"17|{name}".encode())) + return torch.randint(0, VOCAB, (length,), generator=g, device="cuda").tolist() + + +_CACHE: dict = {} + + +@pytest.fixture(scope="module", autouse=True) +def device_memory(): + """Release the cached model before the next module runs.""" + yield + _CACHE.clear() + gc.collect() + torch.cuda.empty_cache() + + +def loaded() -> model_module.BeamSmallModelPrefillBf16Sm103Tp1: + """The seed-17 model; its generated state is dropped once loaded.""" + if "model" not in _CACHE: + state = generate_state(17) + model = model_module.BeamSmallModelPrefillBf16Sm103Tp1(model_config(), device="cuda") + model.load_weights(state) + del state + torch.cuda.empty_cache() + _CACHE["model"] = model + return _CACHE["model"] + + +@contextlib.contextmanager +def counting_attention_calls(): + """Count thop attention calls (and their row counts) while letting them run.""" + cls = type(thop_attention) + original = cls.__call__ + calls: list = [] + + def counted(self, *args, **kwargs): + calls.append( + (kwargs["local_layer_idx"], kwargs["q"].shape[0], kwargs["q"].detach().clone()) + ) + return original(self, *args, **kwargs) + + cls.__call__ = counted + try: + yield calls + finally: + cls.__call__ = original + + +def read_slot(mgr, request_id: int, layer: int, position: int): + """(K [8, 128], V [8, 128]) of one token slot, through the manager's own addressing.""" + pool = int(mgr.impl.get_layer_group_id(model_module.LayerId(mgr.layer_offsets[layer]))) + cache = mgr.kv_cache_map[request_id] + pages = np.frombuffer( + cache.get_base_page_indices(pool), dtype=np.int32, count=int(cache.num_blocks) + ) + buf = mgr.get_buffers(layer, kv_layout="HND") + row = int(pages[position // TPB]) * mgr.get_layer_page_index_scale(layer) // 2 + return buf[row, 0, :, position % TPB, :], buf[row, 1, :, position % TPB, :] + + +def written_bytes(mgr, request_id: int, length: int) -> torch.Tensor: + """Every written K/V slot of a request, all six layers, as one int16 tensor.""" + rows = [] + for layer in range(6): + for p in range(length): + k, v = read_slot(mgr, request_id, layer, p) + rows.append(torch.cat((k, v)).view(torch.int16)) + return torch.stack(rows).cpu() + + +def bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + if a.dtype != b.dtype or a.shape != b.shape: + return False + view = {torch.bfloat16: torch.int16, torch.float32: torch.int32}.get(a.dtype) + return bool(torch.equal(a.contiguous().view(view), b.contiguous().view(view))) + + +# ------------------------------------------------------------------ lifecycle and stopping + + +def test_generation_lifecycle_counters_and_length_stop() -> None: + """create -> prefill -> decode until LENGTH -> release, with counters and call counts.""" + model = loaded() + prompt = token_stream("generation_prompt", 17) + caller = list(prompt) + with BeamSmallModelGenerator(model) as gen, counting_attention_calls() as calls: + state = gen.create_request(prompt, max_new_tokens=5) + assert state.lifecycle is Lifecycle.NEW + first = gen.prefill(state) + assert state.lifecycle is Lifecycle.PREFILLED and state.consumed == 17 + assert first.logits.shape == (VOCAB,) and first.logits.dtype == torch.float32 + assert first.token_id == int(torch.argmax(first.logits)) == state.pending_token + assert [rows for _, rows, _ in calls] == [17] * 6 + steps = [first] + while state.lifecycle is not Lifecycle.FINISHED: + steps.append(gen.decode(state)) + assert state.finish_reason == FinishReason.LENGTH + assert len(state.output_token_ids) == 5 and steps[-1].finish_reason == FinishReason.LENGTH + # The last sampled token is returned but never consumed: P + G - 1 cached tokens. + assert state.consumed == 17 + 5 - 1 + assert [s.position for s in steps] == [16, 17, 18, 19, 20] + assert len(calls) == 6 * 5 and all(rows == 1 for _, rows, _ in calls[6:]) + for step in steps: + assert ( + bool(torch.isfinite(step.logits).all()) and float(step.logits.abs().max()) <= 30.0 + ) + assert step.token_id == int(torch.argmax(step.logits)) + with pytest.raises(ValueError, match="finished"): + gen.decode(state) + assert len(calls) == 30 + gen.release(state) + assert state.lifecycle is Lifecycle.RELEASED + for call in (gen.decode, gen.prefill, gen.release): + with pytest.raises(ValueError, match="released"): + call(state) + assert prompt == caller + + +def test_end_id_stops_before_the_budget() -> None: + """END_ID at the first sampled token after the first that is new to the run: exactly j + tokens, END_ID, and no model call after the j-th; END_ID also wins at the token limit.""" + model = loaded() + prompt = token_stream("generation_prompt", 17) + with BeamSmallModelGenerator(model) as gen: + free = gen.generate(gen.create_request(prompt, max_new_tokens=8, reuse=False)) + assert free.finish_reason == FinishReason.LENGTH and len(free.token_ids) == 8 + tokens = free.token_ids + j = next(j for j in range(2, 9) if tokens[j - 1] not in tokens[: j - 1]) + for state in list(gen.live_requests): + gen.release(state) + for budget in (8, j): + with counting_attention_calls() as calls: + state = gen.create_request( + prompt, max_new_tokens=budget, end_id=tokens[j - 1], reuse=False + ) + out = gen.generate(state) + assert out.token_ids == tokens[:j] and out.finish_reason == FinishReason.END_ID + assert len(calls) == 6 * j and state.consumed == 17 + j - 1 + gen.release(state) + + +@pytest.mark.parametrize( + "case", + [ + "empty", + "token_out_of_range", + "zero_budget", + "over_capacity", + "bad_end_id", + "end_id_no_sampling", + ], +) +def test_create_request_rejects_invalid_requests(case) -> None: + model = loaded() + with BeamSmallModelGenerator(model) as gen: + kwargs = dict(max_new_tokens=2) + prompt = [1, 2, 3] + if case == "empty": + prompt = [] + elif case == "token_out_of_range": + prompt = [VOCAB] + elif case == "zero_budget": + kwargs["max_new_tokens"] = 0 + elif case == "over_capacity": + prompt, kwargs["max_new_tokens"] = [1] * 1000, 25 + elif case == "bad_end_id": + kwargs["end_id"] = -1 + else: + kwargs.update(end_id=5, sample=False) + with pytest.raises(ValueError): + gen.create_request(prompt, **kwargs) + assert gen.live_requests == () + + +def test_handles_are_owned_and_bounded() -> None: + """A third live state, a foreign handle and a decode before prefill are refused.""" + model = loaded() + with BeamSmallModelGenerator(model) as gen, BeamSmallModelGenerator(model) as other: + a = gen.create_request([1, 2, 3], max_new_tokens=1) + b = gen.create_request([4, 5, 6], max_new_tokens=1) + with pytest.raises(ValueError, match="live"): + gen.create_request([7], max_new_tokens=1) + foreign = other.create_request([1, 2, 3], max_new_tokens=1) + with counting_attention_calls() as calls: + with pytest.raises(ValueError, match="another generator"): + gen.prefill(foreign) + with pytest.raises(ValueError, match="new"): + gen.decode(a) + assert calls == [] + for state in (a, b): + gen.release(state) + other.release(foreign) + + +# ------------------------------------------------------------------ cache contents + + +def test_decode_appends_exactly_its_own_keys_and_values() -> None: + """Each explicit token lands at its absolute position in all six layers, as the packed + K/V the step passed to attention, and no earlier slot changes.""" + model = loaded() + prompt = token_stream("append_prompt", 30) + tokens = token_stream("append_tokens", 4) + with BeamSmallModelGenerator(model) as gen: + state = gen.create_request(prompt, max_new_tokens=4, sample=False) + gen.prefill(state) + for token in tokens: + position = state.consumed + before = written_bytes(gen.manager, state.request_id, position) + with counting_attention_calls() as calls: + gen.decode(state, token) + assert [layer for layer, _, _ in calls] == list(range(6)) + for layer, rows, qkv in calls: + assert rows == 1 + k, v = read_slot(gen.manager, state.request_id, layer, position) + assert bitwise_equal(k, qkv[0, 48 * 128 : 56 * 128].view(8, 128)) + assert bitwise_equal(v, qkv[0, 56 * 128 :].view(8, 128)) + assert torch.equal(written_bytes(gen.manager, state.request_id, position), before) + assert state.consumed == position + 1 + record = state.ledger[-1] + assert ( + record["positions"] == [position, position] and record["input"] == "teacher_forced" + ) + assert record["capacity"] >= position + 1 and record["history"] == position + 1 + gen.release(state) + + +def test_prefix_reuse_claims_whole_pages_and_computes_the_rest() -> None: + """B claims A's two common pages and computes 17 rows; the claimed slots are A's bytes and + equal what a reuse-disabled control computes; a different first token misses.""" + model = loaded() + common = token_stream("reuse_common", 64) + a_prompt = common + token_stream("reuse_suffix_a", 17) + b_prompt = common + token_stream("reuse_suffix_b", 17) + miss = [(common[0] + VOCAB // 2) % VOCAB] + common[1:] + token_stream("reuse_suffix_b", 17) + with BeamSmallModelGenerator(model) as gen: + a = gen.create_request(a_prompt, max_new_tokens=1, sample=False) + gen.prefill(a) + b = gen.create_request(b_prompt, max_new_tokens=1, sample=False) + with counting_attention_calls() as calls: + gen.prefill(b) + assert b.reused == 64 and b.computed_prompt_tokens == 17 + assert [rows for _, rows, _ in calls] == [17] * 6 + assert list(gen.manager.kv_cache_map[b.request_id].committed_tokens[:64]) == common + claimed = written_bytes(gen.manager, b.request_id, 64) + assert torch.equal(claimed, written_bytes(gen.manager, a.request_id, 64)) + # Layers 0-3 of the common tokens depend on the dense blocks only, whose arithmetic per + # token does not depend on the rest of the prompt: a recomputation reproduces them bitwise. + dense = claimed[: 4 * 64] + gen.release(b) + gen.release(a) + m = gen.create_request(miss, max_new_tokens=1, sample=False) + gen.prefill(m) + assert m.reused == 0 and m.computed_prompt_tokens == 81 + gen.release(m) + with BeamSmallModelGenerator(model) as gen: + control = gen.create_request(b_prompt, max_new_tokens=1, sample=False, reuse=False) + gen.prefill(control) + assert control.reused == 0 and control.computed_prompt_tokens == 81 + assert torch.equal(written_bytes(gen.manager, control.request_id, 64)[: 4 * 64], dense) + gen.release(control) + + +def test_two_live_requests_stay_isolated() -> None: + """A and B share a reused prefix and advance alternately; neither step changes any slot the + other has written, shared or private.""" + model = loaded() + common = token_stream("reuse_common", 64) + with BeamSmallModelGenerator(model) as gen: + a = gen.create_request(common + token_stream("iso_a", 17), max_new_tokens=3, sample=False) + b = gen.create_request(common + token_stream("iso_b", 17), max_new_tokens=3, sample=False) + gen.prefill(a) + gen.prefill(b) + assert b.reused == 64 + for token_a, token_b in zip(token_stream("iso_ta", 3), token_stream("iso_tb", 3)): + for state, other, token in ((a, b, token_a), (b, a, token_b)): + before = written_bytes(gen.manager, other.request_id, other.consumed) + gen.decode(state, token) + assert torch.equal( + written_bytes(gen.manager, other.request_id, other.consumed), before + ) + gen.release(a) + gen.release(b) + + +def test_weights_cannot_change_under_an_open_cache() -> None: + """load_weights is refused while a generator is open. A parameter replaced by another tensor + (``load_state_dict(assign=True)``) or changed in place is detected before the next step: a + prefill that would reuse a released request's prefix claims nothing and calls no attention, + and a live request's decode allocates nothing. Restoring the bound parameter makes the cache + usable again; after close a new generator starts from an empty cache.""" + model = loaded() + tensors = model.checkpoint_tensors() + original = model.layers[0].self_attn_k_proj + assert tensors["model.layers.0.self_attn.k_proj.weight"] is original + prompt = token_stream("reuse_common", 64) + token_stream("reuse_suffix_a", 17) + + def assign_k_proj(tensor: torch.Tensor) -> None: + model.load_state_dict({"layers.0.self_attn_k_proj": tensor}, strict=False, assign=True) + + with BeamSmallModelGenerator(model) as gen: + released = gen.create_request(prompt, max_new_tokens=1, sample=False) + gen.prefill(released) + gen.release(released) + live = gen.create_request( + token_stream("generation_prompt", 17), max_new_tokens=2, sample=False + ) + gen.prefill(live) + versions = {n: t._version for n, t in tensors.items()} + with pytest.raises(RuntimeError, match="open generator"): + model.load_weights(tensors) + assert {n: t._version for n, t in tensors.items()} == versions + state = gen.create_request(prompt, max_new_tokens=1, sample=False) + capacity = gen.manager.kv_cache_map[live.request_id].capacity + try: + assign_k_proj(torch.zeros_like(original)) + assert model.layers[0].self_attn_k_proj is not original + assert original._version == versions["model.layers.0.self_attn.k_proj.weight"] + with counting_attention_calls() as calls: + with pytest.raises(ValueError, match="changed"): + gen.prefill(state) + with pytest.raises(ValueError, match="changed"): + gen.decode(live, 7) + assert calls == [] and state.request_id not in gen.manager.kv_cache_map + assert state.lifecycle is Lifecycle.NEW and state.reused == 0 + assert live.lifecycle is Lifecycle.PREFILLED and live.consumed == 17 + assert gen.manager.kv_cache_map[live.request_id].capacity == capacity + finally: + assign_k_proj(original) + assert model.layers[0].self_attn_k_proj is original + gen.prefill(state) + assert state.reused == 64 + gen.decode(live, 7) + with torch.no_grad(): + original.mul_(1.0) + with counting_attention_calls() as calls, pytest.raises(ValueError, match="changed"): + gen.decode(live, 7) + assert calls == [] + model.load_weights(tensors) + with BeamSmallModelGenerator(model) as gen: + state = gen.create_request(prompt, max_new_tokens=1, sample=False) + gen.prefill(state) + assert state.reused == 0 + gen.release(state) + + +# ------------------------------------------------------------------ guards before dispatch + + +@contextlib.contextmanager +def step_metadata(gen, state, num_tokens: int, past: int, context: bool, **overrides): + md = TrtllmAttentionMetadata( + max_num_requests=1, + max_num_tokens=model_module.MAX_TOKENS, + kv_cache_manager=gen.manager, + num_contexts=1 if context else 0, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[past]), + seq_lens=torch.tensor([num_tokens], dtype=torch.int), + request_ids=[state.request_id], + prompt_lens=[num_tokens if context else state.prompt_len], + runtime_features=overrides.pop( + "runtime_features", AttentionRuntimeFeatures(cache_reuse=True) + ), + ) + md.prepare() + yield md + + +_BAD_STEPS = [ + "positions_restart", + "two_token_generation", + "cached_context_without_paged_fmha", + "unknown_request", + "beyond_capacity", + "block_offset_corrupted", + "rewound_into_history", + "gap_after_history", +] + + +@pytest.mark.parametrize("case", _BAD_STEPS) +def test_model_rejects_inconsistent_steps_before_dispatch(case) -> None: + """A step whose positions, metadata or block table do not describe the request's own cache + is refused with ValueError before the embedding or any attention call, and no written slot + changes. That includes a generation step that does not begin where the cache's history ends + -- one token back, into the cached history, or one token past it -- although the pages and + capacity it addresses are valid.""" + model = loaded() + prompt = token_stream("guard_prompt", 40) + with BeamSmallModelGenerator(model) as gen: + state = gen.create_request(prompt, max_new_tokens=2, sample=False) + gen.prefill(state) + request = state._request + request.add_new_token(7, 0) + # Capacity for the next token: none beyond the history for "beyond_capacity", one more + # token than usual for "gap_after_history", so that only the history check refuses it. + for _ in range({"beyond_capacity": 0, "gap_after_history": 2}.get(case, 1)): + assert gen.manager.try_allocate_generation(request) + past, ids = state.consumed, torch.tensor([[7]], device="cuda") + assert gen.manager.kv_cache_map[state.request_id].history_length == past == 40 + positions = torch.tensor([[past]], device="cuda") + md_args = dict(num_tokens=1, past=past, context=False) + if case == "positions_restart": + positions = torch.tensor([[0]], device="cuda") + elif case == "two_token_generation": + ids, positions = ( + torch.tensor([[7, 8]], device="cuda"), + torch.tensor([[past, past + 1]], device="cuda"), + ) + md_args = dict(num_tokens=2, past=past - 1, context=False) + elif case == "cached_context_without_paged_fmha": + md_args = dict( + num_tokens=1, past=past, context=True, runtime_features=AttentionRuntimeFeatures() + ) + elif case in ("rewound_into_history", "gap_after_history"): + start = past - 1 if case == "rewound_into_history" else past + 1 + md_args = dict(num_tokens=1, past=start, context=False) + positions = torch.tensor([[start]], device="cuda") + written = written_bytes(gen.manager, state.request_id, past) + with step_metadata(gen, state, **md_args) as md: + if case == "block_offset_corrupted": + md.kv_cache_block_offsets[0, 0, 0, 0] += 2 + elif case == "unknown_request": + md.request_ids = [state.request_id + 1000] + with counting_attention_calls() as calls, pytest.raises(ValueError): + model(ids, positions, md) + assert calls == [] + assert torch.equal(written_bytes(gen.manager, state.request_id, past), written) + gen.release(state) + + +def test_model_rejects_a_context_step_rewound_over_a_reused_prefix() -> None: + """B claims A's two committed pages and holds 64 tokens. A context step for B that starts at + token 32 -- inside those shared pages, on valid pages with capacity and positions that match + it -- is refused before any attention call, and the shared pages keep A's bytes.""" + model = loaded() + common = token_stream("reuse_common", 64) + a_prompt = common + token_stream("reuse_suffix_a", 17) + b_prompt = common + token_stream("reuse_suffix_b", 17) + with BeamSmallModelGenerator(model) as gen: + a = gen.create_request(a_prompt, max_new_tokens=1, sample=False) + gen.prefill(a) + b = gen.create_request(b_prompt, max_new_tokens=1, sample=False) + # The steps prefill takes before its model call: claim the prefix, size the cache. + request = b._request + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + gen.seq_slots.prepare_resources(batch) + assert gen.manager.prepare_context(request) and request.context_current_position == 64 + assert gen.manager.resize_context(request, request.context_remaining_length) + cache = gen.manager.kv_cache_map[b.request_id] + assert cache.history_length == 64 and cache.capacity == 81 + shared = written_bytes(gen.manager, a.request_id, 81) + ids = torch.tensor([b_prompt[32:]], device="cuda") + positions = torch.arange(32, 81, device="cuda").view(1, -1) + with step_metadata(gen, b, num_tokens=49, past=32, context=True) as md: + with counting_attention_calls() as calls, pytest.raises(ValueError, match="holds 64"): + model(ids, positions, md) + assert calls == [] + assert torch.equal(written_bytes(gen.manager, a.request_id, 81), shared) + gen.release(b) + gen.release(a) + + +# ------------------------------------------------------------------ sampler and window + + +def test_torch_sampler_greedy_choice_on_identical_logits() -> None: + """The generator's sampler step picks the lowest index among exact maxima, on a context row + and on a generation row.""" + sampler = TorchSampler(generation.SAMPLER_ARGS) + slots = SeqSlotManager(generation.GENERATOR_CACHE["max_batch_size"]) + request = LlmRequest( + request_id=1, + max_new_tokens=4, + input_tokens=[1, 2, 3], + sampling_config=generation.greedy_sampling_config(), + is_streaming=False, + ) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + slots.prepare_resources(batch) + request.context_chunk_size = 3 + rows = [] + tied = torch.zeros(VOCAB, device="cuda") + tied[[90000, 1234, 151935]] = 2.5 + rows.append((tied, 1234)) + distinct = torch.linspace(-30.0, 30.0, VOCAB, device="cuda").flip(0) + distinct[77] = 30.0 + rows.append((distinct, 0)) + for step, (logits, expected) in enumerate(rows): + generation.sample_step(sampler, request, logits, is_context=step == 0) + assert request.get_last_tokens(0) == expected == int(torch.argmax(logits)) + if step == 0: + assert request.state == LlmRequestState.GENERATION_IN_PROGRESS + slots.free_resources(request) + + +def test_sliding_window_excludes_exactly_the_aged_key() -> None: + """Controlled keys through the generator's cache manager and the catalog attention op: at + query position 512 the sliding layer's output ignores key 0 bit for bit and depends on key + 1, while the full layer depends on key 0.""" + mgr = generation.build_cache_manager() + g = torch.Generator(device="cuda").manual_seed(512) + q = torch.randn(1, 48 * 128, generator=g, device="cuda").to(torch.bfloat16) + kv = torch.randn(513, 2 * 8 * 128, generator=g, device="cuda").to(torch.bfloat16) + request_ids = iter(range(100, 200)) + + def attend(layer: int, kv_rows: torch.Tensor) -> torch.Tensor: + rid = next(request_ids) + tokens = [(rid * 7919 + j * 104729) % VOCAB for j in range(512)] + request = LlmRequest( + request_id=rid, + max_new_tokens=1, + input_tokens=tokens, + sampling_config=generation.greedy_sampling_config(), + is_streaming=False, + ) + try: + assert mgr.prepare_context_cache(request, reuse_limit=0) == 0 + assert mgr.prepare_context(request) and mgr.resize_context(request, 512) + window = 512 if layer < 5 else mgr.max_seq_len + outputs = [] + zero_q = torch.zeros(512, 48 * 128, dtype=torch.bfloat16, device="cuda") + for start, count, context, rows in ( + (0, 512, True, torch.cat((zero_q, kv_rows[:512]), 1)), + (512, 1, False, torch.cat((q, kv_rows[512:513]), 1)), + ): + if not context: + request.move_to_next_context_chunk() + request.state = LlmRequestState.GENERATION_IN_PROGRESS + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + mgr.update_context_resources(batch) + request.add_new_token(5, 0) + assert mgr.try_allocate_generation(request) + md = TrtllmAttentionMetadata( + max_num_requests=1, + max_num_tokens=1024, + kv_cache_manager=mgr, + num_contexts=1 if context else 0, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=[start] + ), + seq_lens=torch.tensor([count], dtype=torch.int), + request_ids=[rid], + prompt_lens=[count if context else 512], + runtime_features=AttentionRuntimeFeatures(cache_reuse=True), + ) + md.prepare() + out = torch.empty(count, 48 * 128, dtype=torch.bfloat16, device="cuda") + thop_attention( + q=rows.contiguous(), + output=out, + num_heads=48, + num_kv_heads=8, + head_size=128, + attention_window_size=window, + **_attention.step_args(md, layer), + **_attention.CALL_CONSTANTS, + ) + outputs.append(out) + torch.cuda.synchronize() + return outputs[-1].clone() + finally: + mgr.free_resources(request) + + def perturbed(position: int) -> torch.Tensor: + rows = kv.clone() + rows[position, : 8 * 128] = -rows[position, : 8 * 128] + rows[position, 8 * 128 :] = rows[position, 8 * 128 :] * 4.0 + 1.0 + return rows + + try: + sliding, full = attend(0, kv), attend(5, kv) + assert bitwise_equal(attend(0, perturbed(0)), sliding) + assert not bitwise_equal(attend(0, perturbed(1)), sliding) + assert not bitwise_equal(attend(5, perturbed(0)), full) + finally: + mgr.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_small_model_prefill.py b/tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_small_model_prefill.py new file mode 100644 index 000000000000..dcb192a34dc3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_small_model_prefill.py @@ -0,0 +1,787 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the Beam six-block small model: bf16 / sm_103 / tp1. + +Configuration under test, fixed below and nowhere else: the first six logical layers of the +checkpoint's schedule -- dense sliding-attention blocks 0-2, sparse sliding-attention blocks +3-4 and the sparse full-attention block 5 without position embedding -- at hidden 5120 with +48 q / 8 kv heads of 128, 256 experts of 2560 per sparse block, vocabulary 151936, untied +embedding and head and a final-logit softcap of 30; bf16 weights and activations with fp32 +routers; one sm_103 GPU; eager, batch 1, one fresh context per call on a six-layer +KVCacheManagerV2; generated weights; seeds {17, 41} x lengths {1, 17, 129}. + +What is checked. Construction, the 2394-tensor loading contract and its rejections, the +input and cache contracts (rejected before any op runs), and on every case the composition +itself, observed read-only at the catalog ops of the forward: the six blocks run once each +in order, each on its predecessor's output, with its own residual coefficients, attention +kind, window and cache slot; layer 5's q takes the query scale and its k is the normalized +k unrotated; the final norm, head and softcap run once each, and the returned logits are +exactly the softcap of the head output. Numerical agreement with the source model is +established separately; these tests generate everything at runtime and need no model files. + +Invocation (fp32 steps need the TF32 overrides removed before torch is imported):: + + env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE python -m pytest \\ + tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1 + +Dependencies: a TensorRT-LLM build for sm_103 with flashinfer importable and about 200 GiB +of free device memory: the generated state and the loaded model hold 62 GiB each, and one +test constructs a second model. No model files. +""" + +from __future__ import annotations + +import contextlib +import gc +import math +import types +import zlib +from collections import OrderedDict + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501 + flashinfer_silu_and_mul, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_neox_rope import ( + beam_neox_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_nope_query_scale import ( # noqa: E501 + beam_nope_query_scale, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_output_gate import ( + beam_output_gate, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + thop_attention, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.residual.beam_depth_residual import ( + beam_depth_residual, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_dispatch_combine_bf16__sm_103__tp1.ops.moe.beam_dispatch_combine import ( # noqa: E501 + beam_dispatch_combine_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_expert_bf16__sm_103__tp1.ops.activation.swiglu_fp32 import ( # noqa: E501 + swiglu_fp32, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_router_bf16__sm_103__tp1.ops.moe.beam_grouped_topk_router import ( # noqa: E501 + beam_grouped_topk_router, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.small_model_prefill_bf16__sm_103__tp1 import ( # noqa: E501 + model as model_module, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.small_model_prefill_bf16__sm_103__tp1.model import ( # noqa: E501 + BeamSmallModelPrefillBf16Sm103Tp1, + prefill_context, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.small_model_prefill_bf16__sm_103__tp1.ops.activation.beam_logit_softcap import ( # noqa: E501 + beam_logit_softcap, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping + +assert torch.cuda.is_available(), "the Beam six-block model requires a CUDA device" + +pytestmark = pytest.mark.high_cuda_memory + +# -------------------------------------------------- the fixed configuration + +#: The checkpoint's configuration reduced to its first six layers; only these fields are read. +CONFIG = types.MappingProxyType( + dict( + num_hidden_layers=6, + hidden_size=5120, + intermediate_size=20480, + moe_intermediate_size=2560, + num_attention_heads=48, + num_key_value_heads=8, + head_dim=128, + vocab_size=151936, + rms_norm_eps=1e-6, + hidden_act="silu", + attention_bias=False, + attention_dropout=0.0, + sliding_window=512, + layer_types=("sliding_attention",) * 5 + ("full_attention",), + mlp_layer_types=("dense",) * 3 + ("sparse",) * 3, + rope_parameters={ + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + }, + num_experts=256, + num_experts_per_tok=8, + n_group=8, + topk_group=4, + norm_topk_prob=True, + routed_scaling_factor=2.5, + tie_word_embeddings=False, + final_logit_softcapping=30.0, + max_position_embeddings=1048576, + ) +) +SEEDS = (17, 41) +LENGTHS = (1, 17, 129) +CASES = tuple((seed, length) for seed in SEEDS for length in LENGTHS) +CASE_IDS = [f"seed{seed}-T{length}" for seed, length in CASES] +VOCAB = CONFIG["vocab_size"] +HIDDEN = CONFIG["hidden_size"] + +#: Every catalog entry the forward can call. +OPS = ( + flashinfer_rmsnorm, + cublas_mm, + beam_neox_rope, + beam_nope_query_scale, + thop_attention, + beam_output_gate, + beam_depth_residual, + flashinfer_silu_and_mul, + beam_grouped_topk_router, + beam_dispatch_combine_op, + swiglu_fp32, + beam_logit_softcap, +) + + +def model_config(**overrides) -> types.SimpleNamespace: + fields = dict(CONFIG) + fields["layer_types"] = list(fields["layer_types"]) + fields["mlp_layer_types"] = list(fields["mlp_layer_types"]) + fields["rope_parameters"] = { + k: (dict(v) if v else v) for k, v in fields["rope_parameters"].items() + } + fields.update(overrides) + return types.SimpleNamespace(**fields) + + +def depth_scales(layer_idx: int, is_attention: bool) -> tuple[float, float]: + """The source model's depth residual coefficients: the n-th add of the residual stream + scales the branch by 1 / sqrt(n) and the stream by sqrt(n - 1) / sqrt(n).""" + n = 2 * (layer_idx + 1) + (0 if is_attention else 1) + return 1.0 / math.sqrt(n), math.sqrt(n - 1) / math.sqrt(n) + + +# -------------------------------------------------- generated state and inputs + + +def _half_width(name: str, shape: tuple) -> tuple[float, float]: + """(half width, offset) of the uniform draw for one tensor: unit-variance projections, + embeddings of std 1, norm weights 1 +- 0.4, gate bias std 0.75, correction bias std 0.05.""" + if name == model_module.EMBED: + return math.sqrt(3.0), 0.0 + if name.endswith("norm.weight") or name.endswith("layernorm.weight"): + return 0.4, 1.0 + if name.endswith("gate_bias"): + return 0.75 * math.sqrt(3.0), 0.0 + if name.endswith("e_score_correction_bias"): + return 0.05 * math.sqrt(3.0), 0.0 + return math.sqrt(3.0 / shape[-1]), 0.0 + + +def generate_state(seed: int) -> "OrderedDict[str, torch.Tensor]": + """All 2394 tensors, each drawn from its own generator keyed by seed and name, fp32 then + rounded once to the stored dtype.""" + state = OrderedDict() + schema = model_module.weight_schema(model_config()) + for name, (shape, dtype) in schema.items(): + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(f"{seed}|{name}".encode())) + half, offset = _half_width(name, shape) + t = torch.rand(shape, generator=g, device="cuda", dtype=torch.float32) + state[name] = (t.mul_(2.0 * half).sub_(half).add_(offset)).to(dtype) + return state + + +def generate_ids(seed: int, length: int) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + return torch.randint(0, VOCAB, (1, max(LENGTHS)), generator=g, device="cuda")[:, :length] + + +_CACHE: dict = {} + + +@pytest.fixture(scope="module", autouse=True) +def device_memory(): + """Release the generated weights and model before the next module runs.""" + yield + _CACHE.clear() + gc.collect() + torch.cuda.empty_cache() + + +def loaded(seed: int) -> tuple[BeamSmallModelPrefillBf16Sm103Tp1, dict]: + """The seed's model, loaded with the seed's generated state; one seed at a time.""" + if _CACHE.get("seed") != seed: + _CACHE.clear() + torch.cuda.empty_cache() + state = generate_state(seed) + model = BeamSmallModelPrefillBf16Sm103Tp1(model_config(), device="cuda") + model.load_weights(state) + _CACHE.update(seed=seed, state=state, model=model) + return _CACHE["model"], _CACHE["state"] + + +def positions_of(length: int) -> torch.Tensor: + return torch.arange(length, device="cuda").view(1, length) + + +def run(model, ids: torch.Tensor) -> torch.Tensor: + length = ids.shape[1] + with prefill_context(length) as md: + out = model(ids, positions_of(length), md) + torch.cuda.synchronize() + return out + + +def bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + """Same dtype, shape and bits (so -0.0 differs from +0.0).""" + if a.dtype != b.dtype or a.shape != b.shape: + return False + view = {torch.bfloat16: torch.int16, torch.float32: torch.int32}.get(a.dtype) + if view is None: + return bool(torch.equal(a, b)) + return bool(torch.equal(a.contiguous().view(view), b.contiguous().view(view))) + + +# -------------------------------------------------- op-path instrumentation + + +@contextlib.contextmanager +def recording(model): + """Log every catalog op call of the forward in order, with the block it ran in, its + activation arguments and its output; and every block's entry and exit.""" + calls: list[dict] = [] + events: list = [] + current = [None] + saved = [] + weight_args_from = {"CublasMm": 1, "FlashinferRmsnorm": 1, "BeamGroupedTopkRouter": 1} + for w in OPS: + cls = type(w) + original = cls.__call__ + saved.append((cls, original)) + + def logged(self, *args, _original=original, _name=cls.__name__, **kwargs): + keep = weight_args_from.get(_name, len(args)) + entry = { + "op": _name, + "block": current[0], + "args": [a if j < keep else None for j, a in enumerate(args)], + "kwargs": kwargs, + } + if isinstance(kwargs.get("q"), torch.Tensor): + entry["qkv"] = kwargs["q"].clone() # before the op appends from it + calls.append(entry) + entry["output"] = _original(self, *args, **kwargs) + return entry["output"] + + cls.__call__ = logged + handles = [] + for i, block in enumerate(model.layers): + + def pre(module, args, _i=i): + current[0] = _i + events.append(("enter", _i, args[0])) + + def post(module, args, out, _i=i): + events.append(("exit", _i, out)) + current[0] = None + + handles.append(block.register_forward_pre_hook(pre)) + handles.append(block.register_forward_hook(post)) + try: + yield calls, events + finally: + for h in handles: + h.remove() + for cls, original in reversed(saved): + cls.__call__ = original + + +@contextlib.contextmanager +def no_dispatch(): + """Fail on any op call or embedding lookup: rejected inputs are rejected before either.""" + saved = [] + for w in OPS: + cls = type(w) + saved.append((cls, cls.__call__)) + + def refuse(self, *args, _name=cls.__name__, **kwargs): + raise AssertionError(f"{_name} dispatched before the input was rejected") + + cls.__call__ = refuse + embedding = torch.nn.functional.embedding + + def refuse_embedding(*args, **kwargs): + raise AssertionError("the embedding ran before the input was rejected") + + torch.nn.functional.embedding = refuse_embedding + try: + yield + finally: + torch.nn.functional.embedding = embedding + for cls, original in reversed(saved): + cls.__call__ = original + + +@contextlib.contextmanager +def manager_context(length: int, **overrides): + """A KVCacheManagerV2 other than the model's own, with prepared metadata for one request.""" + windows = overrides.pop("max_attention_window", list(model_module.ATTENTION_WINDOWS)) + args = dict( + num_layers=6, + num_kv_heads=8, + head_dim=128, + tokens_per_block=32, + max_seq_len=1024, + max_batch_size=1, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.BF16, + ) + args.update(overrides) + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + from tensorrt_llm._torch.metadata import KVCacheParams + + mgr = KVCacheManagerV2( + KvCacheConfig(max_tokens=1024, enable_block_reuse=False, max_attention_window=windows), + CacheTypeCpp.SELF, + **args, + ) + try: + mgr.add_dummy_requests([0], token_nums=[length]) + md = TrtllmAttentionMetadata( + max_num_requests=1, + max_num_tokens=model_module.MAX_TOKENS, + kv_cache_manager=mgr, + num_contexts=1, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0]), + seq_lens=torch.tensor([length], dtype=torch.int), + request_ids=[0], + prompt_lens=[length], + ) + md.prepare() + yield md + finally: + mgr.shutdown() + + +# -------------------------------------------------- tests: configuration and construction + + +def test_configuration_is_the_fixed_scope() -> None: + assert model_module.NUM_LAYERS == 6 + assert model_module.LAYER_TYPES == tuple(CONFIG["layer_types"]) + assert model_module.MLP_LAYER_TYPES == tuple(CONFIG["mlp_layer_types"]) + assert model_module.ATTENTION_WINDOWS == (512,) * 5 + (1024,) + assert model_module.MAX_TOKENS == model_module.KV_CACHE["max_seq_len"] == 1024 + for field, value in model_module.SUPPORTED_CONFIG.items(): + assert CONFIG[field] == value, field + schema = model_module.weight_schema(model_config()) + assert len(schema) == 2394 + assert sum(1 for n in schema if ".mlp.experts." in n) == 3 * 256 * 3 + fp32 = sorted(n for n, (_, dtype) in schema.items() if dtype == torch.float32) + assert fp32 == sorted( + f"model.layers.{i}.mlp.gate.{t}" + for i in (3, 4, 5) + for t in ("weight", "e_score_correction_bias") + ) + + +@pytest.mark.parametrize( + "overrides,match", + [ + (dict(num_hidden_layers=52), "num_hidden_layers"), + (dict(tie_word_embeddings=True), "tie_word_embeddings"), + (dict(final_logit_softcapping=None), "final_logit_softcapping"), + (dict(vocab_size=151552), "vocab_size"), + (dict(layer_types=["sliding_attention"] * 6), "layer_types"), + (dict(mlp_layer_types=["dense"] * 4 + ["sparse"] * 2), "mlp_layer_types"), + ( + dict( + rope_parameters={"sliding_attention": {"rope_type": "default", "rope_theta": 1e4}} + ), + "full_attention", + ), + ], +) +def test_rejects_unsupported_configuration(overrides, match) -> None: + with pytest.raises(ValueError, match=match): + BeamSmallModelPrefillBf16Sm103Tp1(model_config(**overrides), device="cuda") + + +def test_constructs_six_distinct_blocks() -> None: + model, _ = loaded(17) + assert [type(b).__name__ for b in model.layers] == ["BeamDenseBlockBf16Sm103Tp1"] * 3 + [ + "BeamSparseBlockBf16Sm103Tp1" + ] * 3 + assert [b.layer_idx for b in model.layers] == list(range(6)) + for i, block in enumerate(model.layers): + assert block.attention_scales == depth_scales(i, True) + assert block.feed_forward_scales == depth_scales(i, False) + assert [getattr(b, "is_sliding", True) for b in model.layers] == [True] * 5 + [False] + assert [b.rope_theta for b in model.layers] == [10000.0] * 5 + [None] + for i in (3, 4, 5): + assert len(model.layers[i].mlp.experts) == 256 + assert model.layers[i].mlp.gate.weight.dtype == torch.float32 + tensors = model.checkpoint_tensors() + assert len(tensors) == 2394 + spans = sorted( + (t.data_ptr(), t.data_ptr() + t.numel() * t.element_size()) for t in tensors.values() + ) + assert all(a[1] <= b[0] for a, b in zip(spans, spans[1:])), "two tensors share storage" + + +def test_unloaded_model_refuses_to_run() -> None: + loaded(17) + fresh = BeamSmallModelPrefillBf16Sm103Tp1(model_config(), device="cuda") + with prefill_context(1) as md, no_dispatch(), pytest.raises(ValueError, match="load_weights"): + fresh(generate_ids(17, 1), positions_of(1), md) + del fresh + torch.cuda.empty_cache() + + +# -------------------------------------------------- tests: loading + + +def test_load_weights_maps_every_tensor() -> None: + """Every tensor lands under its own name, value, dtype and [out, in] orientation.""" + model, state = loaded(17) + tensors = model.checkpoint_tensors() + assert list(tensors) == list(state) + for name, value in state.items(): + assert bitwise_equal(tensors[name], value), name + assert tensors[model_module.EMBED].data_ptr() != tensors[model_module.HEAD].data_ptr() + assert not torch.equal(tensors[model_module.EMBED], tensors[model_module.HEAD]) + + +_BAD_LOADS = ["missing_head", "unexpected_name", "late_expert_fp32", "late_expert_transposed"] +_BAD_LOADS += ["router_bf16", "head_not_a_tensor"] + + +@pytest.mark.parametrize("case", _BAD_LOADS) +def test_rejected_load_changes_nothing(case) -> None: + """A malformed mapping is rejected before any tensor is copied, even when only its last + tensors are wrong, and the model keeps running on what it held.""" + model, state = loaded(17) + before = run(model, generate_ids(17, 1)) + versions = {n: t._version for n, t in model.checkpoint_tensors().items()} + bad = OrderedDict(state) + late = "model.layers.5.mlp.experts.255.down_proj.weight" + if case == "missing_head": + del bad[model_module.HEAD] + elif case == "unexpected_name": + bad["model.layers.6.input_layernorm.weight"] = state[model_module.NORM] + elif case == "late_expert_fp32": + bad[late] = state[late].float() + elif case == "late_expert_transposed": + bad[late] = state[late].t() + elif case == "router_bf16": + bad["model.layers.5.mlp.gate.weight"] = state["model.layers.5.mlp.gate.weight"].bfloat16() + else: + bad[model_module.HEAD] = state[model_module.HEAD].tolist + with pytest.raises((ValueError, TypeError)): + model.load_weights(bad) + after = {n: t._version for n, t in model.checkpoint_tensors().items()} + assert after == versions + assert bitwise_equal(run(model, generate_ids(17, 1)), before) + + +def test_failed_copy_leaves_the_model_unavailable(monkeypatch) -> None: + model, state = loaded(17) + before = run(model, generate_ids(17, 17)) + + def broken(weights): + raise RuntimeError("copy failed") + + monkeypatch.setattr(model.layers[4], "load_weights", broken) + with pytest.raises(RuntimeError, match="copy failed"): + model.load_weights(state) + with prefill_context(1) as md, no_dispatch(), pytest.raises(ValueError, match="load_weights"): + model(generate_ids(17, 1), positions_of(1), md) + monkeypatch.undo() + model.load_weights(state) + assert bitwise_equal(run(model, generate_ids(17, 17)), before) + + +# -------------------------------------------------- tests: input and cache contracts + +_BAD_INPUTS = [ + "float_ids", + "bool_ids", + "int16_ids", + "negative_id", + "id_vocab_size", + "batch_2", + "flat_ids", + "empty", + "too_long", + "host_ids", + "non_contiguous_ids", + "positions_offset", + "float_positions", + "positions_shape", + "host_positions", + "metadata_type", + "cached_prefix", + "context_length_mismatch", + "one_layer_cache", + "seven_layer_cache", + "duplicate_slots", + "slot_out_of_range", + "missing_layer", + "all_sliding_windows", + "no_sliding_windows", + "fp16_cache", + "four_kv_heads", +] +#: Host pool tables replaced after the manager built them: rows sharing a slot's storage or +#: swapped between slots, an offset outside its pool, a fractional, int64 or misshapen +#: mapping, and pointer rows other than the manager's. +_BAD_POOL_TABLES = [ + "aliased_pool_rows", + "swapped_pool_offsets", + "pool_offset_out_of_range", + "fractional_pool_offsets", + "int64_pool_mapping", + "extra_pool_mapping_row", + "aliased_pool_pointers", + "secondary_pool_pointer", +] +_BAD_INPUTS += _BAD_POOL_TABLES + ["swa_scratch_cache"] + + +@pytest.mark.parametrize("case", _BAD_INPUTS) +def test_rejects_unsupported_inputs_before_dispatch(case, monkeypatch) -> None: + """Each declared restriction raises ValueError before any op or the embedding runs.""" + model, _ = loaded(17) + length = 17 + ids = generate_ids(17, length) + positions = positions_of(length) + context = lambda: prefill_context(length) # noqa: E731 + if case == "float_ids": + ids = ids.float() + elif case == "bool_ids": + ids = ids > 0 + elif case == "int16_ids": + ids = ids.to(torch.int16) + elif case == "negative_id": + ids = ids.clone() + ids[0, 5] = -1 + elif case == "id_vocab_size": + ids = ids.clone() + ids[0, -1] = VOCAB + elif case == "batch_2": + ids = torch.cat((ids, ids)) + positions = torch.cat((positions, positions)) + elif case == "flat_ids": + ids = ids.view(length) + elif case == "empty": + ids, positions = ids[:, :0], positions[:, :0] + elif case == "too_long": + ids = generate_ids(17, 129).repeat(1, 8)[:, : model_module.MAX_TOKENS + 1].contiguous() + positions = positions_of(model_module.MAX_TOKENS + 1) + elif case == "host_ids": + ids = ids.cpu() + elif case == "non_contiguous_ids": + ids = generate_ids(17, 129).repeat(1, 2)[:, : 2 * length : 2] + elif case == "positions_offset": + positions = positions + 1 + elif case == "float_positions": + positions = positions.float() + elif case == "positions_shape": + positions = positions.view(length) + elif case == "host_positions": + positions = positions.cpu() + elif case == "cached_prefix": + context = lambda: manager_context(length, max_seq_len=1024) # noqa: E731 + elif case == "context_length_mismatch": + context = lambda: prefill_context(length - 1) # noqa: E731 + elif case == "one_layer_cache": + context = lambda: manager_context(length, num_layers=1, max_attention_window=[512]) # noqa: E731 + elif case == "seven_layer_cache": + context = lambda: manager_context( # noqa: E731 + length, num_layers=7, max_attention_window=[512] * 5 + [1024, 512] + ) + elif case == "all_sliding_windows": + context = lambda: manager_context(length, max_attention_window=[512] * 6) # noqa: E731 + elif case == "no_sliding_windows": + context = lambda: manager_context(length, max_attention_window=[1024] * 6) # noqa: E731 + elif case == "fp16_cache": + context = lambda: manager_context(length, dtype=DataType.HALF) # noqa: E731 + elif case == "four_kv_heads": + context = lambda: manager_context(length, num_kv_heads=4) # noqa: E731 + with context() as md: + if case == "metadata_type": + md = object() + elif case == "cached_prefix": + md.kv_lens_runtime[0] = length + 8 + elif case in ("duplicate_slots", "slot_out_of_range", "missing_layer"): + offsets = dict(md.kv_cache_manager.layer_offsets) + if case == "duplicate_slots": + offsets[5] = offsets[4] + elif case == "slot_out_of_range": + offsets[5] = 6 + else: + del offsets[5] + monkeypatch.setattr(md.kv_cache_manager, "layer_offsets", offsets) + elif case in _BAD_POOL_TABLES: + mgr = md.kv_cache_manager + mapping = mgr.kv_cache_pool_mapping.clone() + pointers = mgr.kv_cache_pool_pointers.clone() + if case == "aliased_pool_rows": + mapping[4] = mapping[3] + elif case == "swapped_pool_offsets": + mapping[[0, 1]] = mapping[[1, 0]] + elif case == "pool_offset_out_of_range": + mapping[5, 1] = 1 << 20 + elif case == "fractional_pool_offsets": + mapping = mapping.float() + mapping[5, 1] = 0.5 + elif case == "int64_pool_mapping": + mapping = mapping.long() + elif case == "extra_pool_mapping_row": + mapping = torch.cat((mapping, mapping[5:])) + elif case == "aliased_pool_pointers": + pointers[1, 0] = pointers[0, 0] + else: + pointers[0, 1] = pointers[0, 0] + monkeypatch.setattr(mgr, "kv_cache_pool_mapping", mapping) + monkeypatch.setattr(mgr, "kv_cache_pool_pointers", pointers) + elif case == "swa_scratch_cache": + monkeypatch.setattr(md.kv_cache_manager, "enable_swa_scratch_reuse", True) + with no_dispatch(), pytest.raises(ValueError): + model(ids, positions, md) + + +def test_cache_slots_are_the_managers_own_buffers() -> None: + """The prepared cache passes check_cache, and each layer's pool row addresses the buffer + the manager itself reports for that layer: six distinct buffers, as allocated.""" + with prefill_context(17) as md: + model_module.check_cache(md, 17) + mgr = md.kv_cache_manager + mapping, pointers = md.host_kv_cache_pool_mapping, md.host_kv_cache_pool_pointers + assert mapping.dtype == torch.int32 and pointers.dtype == torch.int64 + page = 32 * 8 * 128 * 2 + addresses = [] + for layer_idx in range(6): + pool, offset = mapping[mgr.layer_offsets[layer_idx]].tolist() + addresses.append(int(pointers[pool, 0]) + offset * 2 * page) + assert addresses[-1] == mgr.get_buffers(layer_idx).data_ptr(), layer_idx + assert len(set(addresses)) == 6 + + +# -------------------------------------------------- tests: the composed forward + + +@pytest.mark.parametrize("seed,length", CASES, ids=CASE_IDS) +def test_forward_composes_the_six_blocks(seed, length) -> None: + """One plain call and one observed call: the same logits, every element of them, and the + composition the observation shows.""" + model, _ = loaded(seed) + ids = generate_ids(seed, length) + ids_before = ids.clone() + versions = {n: t._version for n, t in model.checkpoint_tensors().items()} + plain = run(model, ids) + assert plain.shape == (1, length, VOCAB) and plain.dtype == torch.float32 + assert plain.device == model.device and bool(torch.isfinite(plain).all()) + assert float(plain.abs().max()) <= 30.0 + with prefill_context(length) as md, recording(model) as (calls, events): + observed = model(ids, positions_of(length), md) + torch.cuda.synchronize() + assert bitwise_equal(observed, plain) + assert torch.equal(ids, ids_before) + assert {n: t._version for n, t in model.checkpoint_tensors().items()} == versions + # Blocks 0..5 once each, in order, each on its predecessor's output. + assert [(kind, i) for kind, i, _ in events] == [ + (kind, i) for i in range(6) for kind in ("enter", "exit") + ] + for i in range(1, 6): + assert bitwise_equal(events[2 * i][2], events[2 * i - 1][2]) + embedded = events[0][2] + assert embedded.shape == (1, length, HIDDEN) and embedded.dtype == torch.bfloat16 + assert bitwise_equal(embedded[0], model.embed_tokens[ids[0]]) + for i in range(6): + block = [c for c in calls if c["block"] == i and c["op"] != "expert"] + attention = [c for c in block if c["op"] == "ThopAttention"] + (att,) = attention + assert att["kwargs"]["local_layer_idx"] == i + assert att["kwargs"]["attention_window_size"] == (512 if i < 5 else 1024) + residuals = [c for c in block if c["op"] == "BeamDepthResidual"] + assert [tuple(c["args"][2:4]) for c in residuals] == [ + depth_scales(i, True), + depth_scales(i, False), + ] + qkv = att["qkv"] + k_norm = [c for c in block if c["op"] == "FlashinferRmsnorm"][2]["output"] + q_norm = [c for c in block if c["op"] == "FlashinferRmsnorm"][1]["output"] + packed_q = qkv[:, : 48 * 128].reshape(length, 48, 128) + packed_k = qkv[:, 48 * 128 : 56 * 128].reshape(length, 8, 128) + if i < 5: + assert [c["op"] for c in block].count("BeamNeoxRope") == 1 + assert "BeamNopeQueryScale" not in [c["op"] for c in block] + else: + assert "BeamNeoxRope" not in [c["op"] for c in block] + assert bitwise_equal(packed_k, k_norm.view(length, 8, 128)) + assert bitwise_equal( + packed_q, + beam_nope_query_scale(q_norm.view(length, 48, 128), positions_of(length)[0]), + ) + top = [c for c in calls if c["block"] is None] + assert [c["op"] for c in top] == ["FlashinferRmsnorm", "CublasMm", "BeamLogitSoftcap"] + norm, head, softcap = top + assert bitwise_equal(norm["args"][0], events[-1][2].view(length, HIDDEN)) + assert bitwise_equal(head["args"][0], norm["output"]) + assert head["output"].dtype == torch.bfloat16 and head["output"].shape == (length, VOCAB) + # The softcap, exactly once: the returned logits are this expression of the head output. + once = (head["output"].float() / 30.0).tanh() * 30.0 + assert bitwise_equal(plain.view(length, VOCAB), once) + assert bitwise_equal(softcap["args"][0], head["output"]) and softcap["args"][1] == 30.0 + + +def test_layer5_query_scale_is_visible() -> None: + """At 129 tokens the query scale changes block 5's q at later positions, so an omitted + scale would be a different call.""" + model, _ = loaded(17) + with prefill_context(129) as md, recording(model) as (calls, _): + model(generate_ids(17, 129), positions_of(129), md) + torch.cuda.synchronize() + block = [c for c in calls if c["block"] == 5] + q_norm = [c for c in block if c["op"] == "FlashinferRmsnorm"][1]["output"].view(129, 48, 128) + (att,) = [c for c in block if c["op"] == "ThopAttention"] + packed_q = att["qkv"][:, : 48 * 128].reshape(129, 48, 128) + changed = (packed_q != q_norm).flatten(1).any(dim=1) + assert not bool(changed[0]) and bool(changed[-1]) + + +def test_causal_prefix_and_request_isolation() -> None: + """Changing the tokens after the first 17 leaves the dense blocks' prefix rows and block + 3's input to its feed-forward unchanged, bit for bit; and a request between two identical + requests changes nothing of the second.""" + model, _ = loaded(17) + full = generate_ids(17, 129) + other = full.clone() + other[:, 17:] = generate_ids(41, 129)[:, 17:] + + def observe(ids): + with prefill_context(ids.shape[1]) as md, recording(model) as (calls, events): + out = model(ids, positions_of(ids.shape[1]), md) + torch.cuda.synchronize() + pre_ff = [c["output"] for c in calls if c["block"] == 3 and c["op"] == "FlashinferRmsnorm"][ + 4 + ] + return out, [x for kind, _, x in events if kind == "exit"], pre_ff + + out_a, exits_a, pre_ff_a = observe(full) + out_b, exits_b, pre_ff_b = observe(other) + for i in range(3): + assert bitwise_equal(exits_a[i][:, :17], exits_b[i][:, :17]) + assert bitwise_equal(pre_ff_a[:17], pre_ff_b[:17]) + assert not torch.equal(out_a[:, 17:], out_b[:, 17:]) + first = run(model, full[:, :17]) + run(model, generate_ids(41, 17)) + assert bitwise_equal(run(model, full[:, :17]), first) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/_sparse_block_reference.py b/tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/_sparse_block_reference.py new file mode 100644 index 000000000000..9abed732ddff --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/_sparse_block_reference.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plain-Torch reference for the sparse sliding-attention decoder block 3 of Beam. + +An independent restatement of the source block for generated-input comparison: sandwich +RMSNorms (fp32 statistics and affine, one bf16 rounding); separate bias-free q/k/v/gate +projections in [out, in] layout; per-head q/k RMSNorm rounded to bf16 before an fp32 NeoX +RoPE rounded once; causal GQA attention, 6 q heads per kv head, scale 1/sqrt(128), window +512; the per-head sigmoid output gate in bf16 steps; the output projection; the MoE of +``../moe_bf16__sm_103__tp1/_moe_reference.py``; and block 3's depth-scaled residual adds, +(1/sqrt(8), sqrt(7)/sqrt(8)) after attention and (1/3, sqrt(8)/3) after the feed-forward, +each scaled term and each sum rounded to bf16. + +Only torch and that MoE reference are imported. The attention core is torch's SDPA, which +the component never runs. State and input are generated on the GPU as the MoE reference +generates its own, one seeded generator per tensor; the block's 770 ``mlp.`` tensors are +exactly the MoE reference's state of the same seed. +""" + +from __future__ import annotations + +import math +from collections import OrderedDict + +import torch +import torch.nn.functional as F + +__extra_import_path__ = ["../moe_bf16__sm_103__tp1"] +import _moe_reference as moe_reference # noqa: E402 + +HIDDEN = 5120 +HEADS = 48 +KV_HEADS = 8 +HEAD_DIM = 128 +EPS = 1e-6 +ROPE_THETA = 10000.0 +WINDOW = 512 +LAYER_IDX = 3 +MAX_TOKENS = 129 +SEEDS = (17, 41) +LENGTHS = (1, 17, 129) +#: The MoE's tensors carry this prefix among the block's names. +MLP_PREFIX = "mlp." + +STATE_RECIPE = { + "seeding": "one torch.Generator on the device per call, re-seeded for every tensor with " + "zlib.crc32(f'{seed}|{name}'); the mlp. tensors are the MoE reference's state of the " + "seed, drawn under their names relative to mlp.", + "*norm.weight": "1 + 0.25 * N(0, 1), drawn in fp32 and rounded once to bf16", + "self_attn.gate_bias": "0.75 * N(0, 1), drawn in fp32 and rounded once to bf16", + "projections": "N(0, 1) / sqrt(input width), drawn in fp32 and rounded once to bf16", + "input": "N(0, 1) of shape [1, 129, 5120] under the name 'block_input', drawn in fp32 " + "and rounded once to bf16; a T-token case uses the first T tokens, positions arange(T)", +} + + +def depth_coefs(layer_idx: int, is_attention: bool) -> tuple[float, float]: + """(branch, stream) scales of the n-th residual add, 1 / sqrt(n) and sqrt(n - 1) / + sqrt(n); block i's attention add is the 2(i + 1)-th, its feed-forward add the next.""" + num_adds = 2 * (layer_idx + 1) + (0 if is_attention else 1) + denom = math.sqrt(num_adds) + return 1.0 / denom, math.sqrt(num_adds - 1) / denom + + +ATTENTION_SCALES = depth_coefs(LAYER_IDX, True) +FEED_FORWARD_SCALES = depth_coefs(LAYER_IDX, False) + + +def _ignore(name: str, value: object) -> None: + del name, value + + +# ----------------------------------------------------------------------------- state and inputs + + +def own_shapes() -> "OrderedDict[str, tuple[int, ...]]": + """The block's twelve bf16 attention and norm tensors by checkpoint name, [out, in].""" + q, kv = HEADS * HEAD_DIM, KV_HEADS * HEAD_DIM + return OrderedDict( + [ + ("input_layernorm.weight", (HIDDEN,)), + ("self_attn.gate_bias", (HEADS,)), + ("self_attn.q_proj.weight", (q, HIDDEN)), + ("self_attn.k_proj.weight", (kv, HIDDEN)), + ("self_attn.v_proj.weight", (kv, HIDDEN)), + ("self_attn.gate_proj.weight", (q, HIDDEN)), + ("self_attn.o_proj.weight", (HIDDEN, q)), + ("self_attn.q_norm.weight", (HEAD_DIM,)), + ("self_attn.k_norm.weight", (HEAD_DIM,)), + ("pre_feedforward_layernorm.weight", (HIDDEN,)), + ("post_attention_layernorm.weight", (HIDDEN,)), + ("post_feedforward_layernorm.weight", (HIDDEN,)), + ] + ) + + +def param_shapes() -> "OrderedDict[str, tuple[tuple[int, ...], torch.dtype]]": + """The block's 782 tensors by checkpoint name: its own twelve, then the MoE's under + ``mlp.``, with their dtypes.""" + shapes = OrderedDict((name, (shape, torch.bfloat16)) for name, shape in own_shapes().items()) + for name, shape_dtype in moe_reference.param_shapes().items(): + shapes[MLP_PREFIX + name] = shape_dtype + return shapes + + +def generate_state( + seed: int, device: str | torch.device = "cuda" +) -> "OrderedDict[str, torch.Tensor]": + """The seed's 782 tensors on ``device``, in checkpoint order (see `STATE_RECIPE`).""" + gen = torch.Generator(device=device) + state: "OrderedDict[str, torch.Tensor]" = OrderedDict() + for name, shape in own_shapes().items(): + gen.manual_seed(moe_reference.tensor_seed(seed, name)) + draw = torch.randn(shape, generator=gen, dtype=torch.float32, device=device) + if name.endswith("norm.weight"): + value = 1.0 + 0.25 * draw + elif name == "self_attn.gate_bias": + value = 0.75 * draw + else: + value = draw / math.sqrt(shape[1]) + state[name] = value.to(torch.bfloat16).contiguous() + for name, tensor in moe_reference.generate_state(seed, device).items(): + state[MLP_PREFIX + name] = tensor + return state + + +def moe_state(state: "dict[str, torch.Tensor]") -> "OrderedDict[str, torch.Tensor]": + """The MoE's tensors of a block state by name relative to ``mlp.``; nothing is copied.""" + return OrderedDict( + (name[len(MLP_PREFIX) :], tensor) + for name, tensor in state.items() + if name.startswith(MLP_PREFIX) + ) + + +def generate_input(seed: int, device: str | torch.device = "cuda") -> torch.Tensor: + """``[1, 129, 5120]`` bf16: the residual stream entering block 3 (see `STATE_RECIPE`).""" + gen = torch.Generator(device=device) + gen.manual_seed(moe_reference.tensor_seed(seed, "block_input")) + x = torch.randn(1, MAX_TOKENS, HIDDEN, generator=gen, dtype=torch.float32, device=device) + return x.to(torch.bfloat16) + + +def state_problems(state: "dict[str, torch.Tensor]", seed: int) -> list[str]: + """Exact accounting plus the nontrivial-state properties the comparisons rely on.""" + shapes = param_shapes() + if list(state) != list(shapes): + return ["names or order differ from param_shapes()"] + problems = [] + own = own_shapes() + for name, shape in own.items(): + tensor = state[name] + if tuple(tensor.shape) != shape or tensor.dtype != torch.bfloat16: + problems.append(f"{name}: {tuple(tensor.shape)} {tensor.dtype}, expected {shape}") + elif not bool(torch.isfinite(tensor).all()): + problems.append(f"{name}: nonfinite") + if problems: + return problems + names = list(own) + list(moe_reference.param_shapes()) + if len({moe_reference.tensor_seed(seed, name) for name in names}) != len(names): + problems.append("two tensors share a generator seed") + norms = [name for name in own if name.endswith("norm.weight")] + for name in norms: + if bool((state[name] == 1).all()): + problems.append(f"{name}: unit affine") + for i, a in enumerate(norms): + for b in norms[i + 1 :]: + if state[a].shape == state[b].shape and torch.equal(state[a], state[b]): + problems.append(f"{a} and {b} are identical") + bias = state["self_attn.gate_bias"] + if bool((bias == bias[0]).all()): + problems.append("self_attn.gate_bias does not vary across heads") + problems.extend(f"MoE: {p}" for p in moe_reference.state_problems(moe_state(state), seed)) + return problems + + +# ----------------------------------------------------------------------------- the block + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Over the last dim: fp32 statistics, normalization and affine, rounded once.""" + x32 = x.float() + normed = x32 * torch.rsqrt(x32.pow(2).mean(dim=-1, keepdim=True) + EPS) + return (normed * weight.float()).to(x.dtype) + + +def rope( + q: torch.Tensor, k: torch.Tensor, positions: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """NeoX half rotation of ``q [T, 48, 128]`` and ``k [T, 8, 128]`` by their token's + position: fp32 angles ``position / 10000^(2i / 128)``, ``x cos + rotate_half(x) sin`` + in fp32, rounded once.""" + exponent = torch.arange(0, HEAD_DIM, 2, dtype=torch.float32, device=positions.device) + inv_freq = 1.0 / (ROPE_THETA ** (exponent / HEAD_DIM)) + angles = positions.float().unsqueeze(-1) * inv_freq + emb = torch.cat((angles, angles), dim=-1) + cos, sin = emb.cos().unsqueeze(1), emb.sin().unsqueeze(1) + half = HEAD_DIM // 2 + + def rotate(x: torch.Tensor) -> torch.Tensor: + x32 = x.float() + rotated = torch.cat((-x32[..., half:], x32[..., :half]), dim=-1) + return (x32 * cos + rotated * sin).to(x.dtype) + + return rotate(q), rotate(k) + + +def attention_core(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """Causal windowed GQA over one fresh context: ``q [T, 48, 128]``, ``k`` and ``v`` + ``[T, 8, 128]`` bf16, q head h reading kv head h // 6; returns ``[T, 48, 128]`` bf16.""" + tokens = q.shape[0] + repeat = HEADS // KV_HEADS + k = k.repeat_interleave(repeat, dim=1) + v = v.repeat_interleave(repeat, dim=1) + position = torch.arange(tokens, device=q.device) + distance = position.unsqueeze(1) - position.unsqueeze(0) + keep = (distance >= 0) & (distance < WINDOW) + out = F.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + k.transpose(0, 1).unsqueeze(0), + v.transpose(0, 1).unsqueeze(0), + attn_mask=keep, + dropout_p=0.0, + scale=HEAD_DIM**-0.5, + ) + return out.squeeze(0).transpose(0, 1) + + +def output_gate(core: torch.Tensor, gate: torch.Tensor, gate_bias: torch.Tensor) -> torch.Tensor: + """``core * sigmoid(gate + gate_bias)`` per head, ``[T, 48, 128]``: the add, the + sigmoid and the product each rounded to bf16.""" + biased = (gate.float() + gate_bias.float().view(1, HEADS, 1)).to(torch.bfloat16) + weight = torch.sigmoid(biased.float()).to(torch.bfloat16) + return (core.float() * weight.float()).to(torch.bfloat16) + + +def attention( + h: torch.Tensor, positions: torch.Tensor, state: "dict[str, torch.Tensor]", record +) -> torch.Tensor: + """Normalized input ``[T, 5120]`` -> attention output ``[T, 5120]``, before the + post-attention norm.""" + tokens = h.shape[0] + q = F.linear(h, state["self_attn.q_proj.weight"]) + k = F.linear(h, state["self_attn.k_proj.weight"]) + v = F.linear(h, state["self_attn.v_proj.weight"]) + gate = F.linear(h, state["self_attn.gate_proj.weight"]) + for name, value in (("q_proj", q), ("k_proj", k), ("v_proj", v), ("gate_proj", gate)): + record(name, value) + q = rms_norm(q.view(tokens, HEADS, HEAD_DIM), state["self_attn.q_norm.weight"]) + k = rms_norm(k.view(tokens, KV_HEADS, HEAD_DIM), state["self_attn.k_norm.weight"]) + record("q_norm", q) + record("k_norm", k) + q, k = rope(q, k, positions) + record("q_rope", q) + record("k_rope", k) + core = attention_core(q, k, v.view(tokens, KV_HEADS, HEAD_DIM)) + record("attention_core", core) + gated = output_gate(core, gate.view(tokens, HEADS, HEAD_DIM), state["self_attn.gate_bias"]) + record("gated", gated) + out = F.linear(gated.reshape(tokens, HEADS * HEAD_DIM), state["self_attn.o_proj.weight"]) + record("attention_output", out) + return out + + +def residual_add( + branch: torch.Tensor, stream: torch.Tensor, scales: tuple[float, float] +) -> torch.Tensor: + """``branch * scales[0] + stream * scales[1]``: each product and the sum computed in + fp32 and rounded to bf16 in turn.""" + branch_scale, stream_scale = scales + scaled_branch = (branch.float() * branch_scale).to(torch.bfloat16) + scaled_stream = (stream.float() * stream_scale).to(torch.bfloat16) + return (scaled_branch.float() + scaled_stream.float()).to(torch.bfloat16) + + +def block( + x: torch.Tensor, + positions: torch.Tensor, + state: "dict[str, torch.Tensor]", + record=None, +) -> torch.Tensor: + """``x [1, T, 5120]`` bf16 and integer ``positions [1, T]`` -> ``[1, T, 5120]`` bf16. + + ``record(name, value)``, when given, receives every intermediate ``[T, ...]``, the + MoE's under ``moe.`` (``moe.routing`` is its route), and ``block_output`` last. + """ + keep = record if record is not None else _ignore + tokens = x.shape[1] + stream = x.reshape(tokens, HIDDEN) + with moe_reference.reference_math(): + h = rms_norm(stream, state["input_layernorm.weight"]) + keep("input_layernorm", h) + a = attention(h, positions.reshape(tokens), state, keep) + a = rms_norm(a, state["post_attention_layernorm.weight"]) + keep("post_attention_layernorm", a) + stream = residual_add(a, stream, ATTENTION_SCALES) + keep("attention_residual", stream) + h = rms_norm(stream, state["pre_feedforward_layernorm.weight"]) + keep("pre_feedforward_layernorm", h) + m = moe_reference.moe(h, moe_state(state), lambda name, value: keep(f"moe.{name}", value)) + keep("moe_output", m) + m = rms_norm(m, state["post_feedforward_layernorm.weight"]) + keep("post_feedforward_layernorm", m) + out = residual_add(m, stream, FEED_FORWARD_SCALES) + keep("block_output", out) + return out.view(1, tokens, HIDDEN) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/test_sparse_block.py b/tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/test_sparse_block.py new file mode 100644 index 000000000000..b1d0fa65ae40 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/test_sparse_block.py @@ -0,0 +1,1453 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the Beam sparse decoder block 3 component: bf16 / sm_103 / tp1. + +Configuration under test, fixed below and nowhere else: block 3 of the 52-block model +(sliding attention, MoE feed-forward) at hidden 5120 with 48 q / 8 kv heads of 128, and +block 3's MoE: 256 experts of intermediate 2560, 8 per token from 4 of 8 groups, routed +scaling factor 2.5; bf16 weights and activations with the source model's fp32 steps and +an fp32 router; one sm_103 GPU; eager, batch 1, one fresh context per call on a +KVCacheManagerV2 whose layer mask selects logical layer 3 alone; generated weights; +seeds {17, 41} x lengths {1, 17, 129}. + +What is compared. ``_sparse_block_reference.py`` restates the source block in plain +torch, with the MoE of ``../moe_bf16__sm_103__tp1/_moe_reference.py``, and generates its +state and input deterministically on the GPU. Routing is judged apart from numerical +error. On the component's own feed-forward input, captured from the forward, the +reference must route exactly as the component did. Against the whole reference block, +whose feed-forward input differs from the component's by rounding, every routing +decision the two take differently must sit on a near tie of the reference's own scores; +the block output is compared on the tokens routed to the same experts, and the MoE +output on the component's own feed-forward input. Each compared tensor is held to one +rule in `RULES`. + +Invocation (fp32 steps need the TF32 overrides removed before torch is imported):: + + env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE python -m pytest \\ + tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1 + +Optional environment: + +* ``BEAM_SPARSE_BLOCK_EVIDENCE``: a directory that receives one JSON record per test + case: every comparison with its rule, the op calls and CUDA kernels of the forward + where they were recorded, and the device and software it ran on; and + ``fp32-session.json``, the fp32 GEMM checks made before and after the module's tests. + +Dependencies: a TensorRT-LLM build for sm_103 with flashinfer importable and about 60 GiB +of free device memory: one seed's generated state and the loaded block hold 18.75 GiB of +expert weights each, and one test builds a second block. No model files. +""" + +from __future__ import annotations + +import contextlib +import gc +import hashlib +import json +import math +import os +import platform +import re +import subprocess +import sys +import types +from collections import OrderedDict +from pathlib import Path + +import pytest +import torch + +import tensorrt_llm +import tensorrt_llm._torch.models.modeling_beam as builtin_beam +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_neox_rope import ( + beam_neox_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_output_gate import ( + beam_output_gate, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + thop_attention, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.residual.beam_depth_residual import ( + beam_depth_residual, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_bf16__sm_103__tp1.moe import ( + BeamMoeBf16Sm103Tp1, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_dispatch_combine_bf16__sm_103__tp1.ops.moe.beam_dispatch_combine import ( # noqa: E501 + beam_dispatch_combine_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_expert_bf16__sm_103__tp1.ops.activation.swiglu_fp32 import ( # noqa: E501 + swiglu_fp32, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.moe_router_bf16__sm_103__tp1.ops.moe.beam_grouped_topk_router import ( # noqa: E501 + beam_grouped_topk_router, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.sparse_block_bf16__sm_103__tp1 import ( # noqa: E501 + block as block_module, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.sparse_block_bf16__sm_103__tp1.block import ( # noqa: E501 + BeamSparseBlockBf16Sm103Tp1, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping + +__extra_import_path__ = [".", "../moe_bf16__sm_103__tp1", "../../.."] +import _moe_reference as moe_reference # noqa: E402 +import _sparse_block_reference as reference # noqa: E402 +from _validating import validating # noqa: E402 + +assert torch.cuda.is_available(), "the Beam sparse block requires a CUDA device" + +pytestmark = pytest.mark.high_cuda_memory + +# -------------------------------------------------- the fixed configuration + +NUM_LAYERS = 52 +LAYER_IDX = 3 +_FULL_ATTENTION_LAYERS = frozenset({5, 11, 17, 23, 29, 35, 41, 47, 51}) + +#: The block as the checkpoint configures it: the 52-block model's layer pattern and +#: block 3's attention and MoE. Only these fields are read. +BLOCK_CONFIG = types.MappingProxyType( + dict( + hidden_size=5120, + num_attention_heads=48, + num_key_value_heads=8, + head_dim=128, + rms_norm_eps=1e-6, + attention_bias=False, + attention_dropout=0.0, + sliding_window=512, + num_hidden_layers=NUM_LAYERS, + layer_types=tuple( + "full_attention" if i in _FULL_ATTENTION_LAYERS else "sliding_attention" + for i in range(NUM_LAYERS) + ), + mlp_layer_types=("dense",) * 3 + ("sparse",) * (NUM_LAYERS - 3), + rope_parameters={ + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + }, + moe_intermediate_size=2560, + num_experts=256, + num_experts_per_tok=8, + n_group=8, + topk_group=4, + norm_topk_prob=True, + routed_scaling_factor=2.5, + hidden_act="silu", + ) +) +DTYPE = torch.bfloat16 +SM = (10, 3) +SEEDS = (17, 41) +LENGTHS = (1, 17, 129) +CASES = tuple((seed, length) for seed in SEEDS for length in LENGTHS) +CASE_IDS = [f"seed{seed}-T{length}" for seed, length in CASES] +#: Eager, batch 1, one fresh context: no CUDA graph, no overlap scheduler, no cached prefix. +EXECUTION = types.MappingProxyType( + dict(eager=True, batch_size=1, fresh_context=True, cuda_graph=False, overlap_scheduler=False) +) +#: The KV cache the attention op appends to: logical layer 3 of 52 selected by the layer +#: mask, bf16, 32-token pages, 512 tokens per sequence (the window), one request. +KV_CACHE = types.MappingProxyType(dict(tokens_per_block=32, max_seq_len=512, num_blocks=16)) +LAYER_MASK = tuple(i == LAYER_IDX for i in range(NUM_LAYERS)) +REQUEST_ID = 0 + +# -------------------------------------------------- comparison rules + +_RTOL = 2.0**-7 + +#: (atol, rtol) per kind of compared bf16 tensor. An element passes when +#: |actual - reference| <= atol + rtol * |reference|, both widened to fp64; every element of +#: every compared tensor must pass, and no NaN or Inf is accepted. rtol 2^-7 is one bf16 unit +#: at any magnitude. "block_output" holds the block output on the tokens routed to the same +#: experts as in the reference block, and the two prefix-identical inputs' common rows; +#: "moe_output" holds the MoE output against the reference MoE on the same input. +RULES = types.MappingProxyType({"block_output": (0.1, _RTOL), "moe_output": (0.04, _RTOL)}) +#: A routing decision the component takes differently from the whole reference block -- a +#: group kept or ordered differently, or an expert selected or ordered differently within +#: the groups both kept -- passes only on a margin of at most this in the reference's own +#: group scores or selection scores. Routing on the same input is never excused: it is +#: compared exactly. +ROUTE_GAP = 3e-3 + +#: The forward's top-level op calls, in order: the composition this component claims to +#: run. The experts' calls nest inside BeamDispatchCombine. +EXPECTED_CALLS = ( + "FlashinferRmsnorm", + "CublasMm", + "CublasMm", + "CublasMm", + "CublasMm", + "FlashinferRmsnorm", + "FlashinferRmsnorm", + "BeamNeoxRope", + "ThopAttention", + "BeamOutputGate", + "CublasMm", + "FlashinferRmsnorm", + "BeamDepthResidual", + "FlashinferRmsnorm", + "BeamGroupedTopkRouter", + "BeamDispatchCombine", + "FlashinferRmsnorm", + "BeamDepthResidual", +) +#: What each selected expert runs, in order, inside BeamDispatchCombine. +EXPERT_CALLS = ("expert", "CublasMm", "CublasMm", "SwigluFp32", "CublasMm") +#: Every catalog entry the block's forward can call. +OPS = ( + flashinfer_rmsnorm, + cublas_mm, + beam_neox_rope, + thop_attention, + beam_output_gate, + beam_depth_residual, + beam_grouped_topk_router, + swiglu_fp32, + beam_dispatch_combine_op, +) +#: Keyword arguments of the entries' calls that identify the path. +_RECORDED_KWARGS = ( + "num_heads", + "num_kv_heads", + "head_size", + "attention_window_size", + "max_seq_len", + "tokens_per_block", + "mask_type", + "is_fused_qkv", + "num_contexts", + "num_ctx_tokens", + "use_paged_context_fmha", + "position_embedding_type", + "local_layer_idx", + "quant_mode", + "n_group", + "topk_group", + "top_k", + "eps", + "routed_scaling_factor", +) + +_SDPA_OPS = re.compile(r"scaled_dot_product|_flash_attention|_efficient_attention|_cudnn_attention") +_ATTENTION_KERNEL = re.compile(r"(?i)fmha|attention") + + +def block_config(**overrides) -> types.SimpleNamespace: + fields = dict(BLOCK_CONFIG) + fields["layer_types"] = list(fields["layer_types"]) + fields["mlp_layer_types"] = list(fields["mlp_layer_types"]) + fields["rope_parameters"] = { + k: (dict(v) if v else v) for k, v in fields["rope_parameters"].items() + } + fields.update(overrides) + return types.SimpleNamespace(**fields) + + +def replaced(items, index: int, value) -> list: + """A list copy of ``items`` with one entry replaced.""" + out = list(items) + out[index] = value + return out + + +# -------------------------------------------------- precision and environment + + +def _fp32_gemm_probe() -> float: + """Worst relative error of an fp32 GEMM against fp64; TF32 substitution shows as ~1e-3.""" + gen = torch.Generator().manual_seed(1234) + x = torch.randn(64, 5120, generator=gen).to(torch.bfloat16).float().cuda() + w = (torch.randn(512, 5120, generator=gen) * 5120**-0.5).float().cuda() + y = torch.nn.functional.linear(x, w) + ref = x.double() @ w.double().t() + return float((y.double() - ref).abs().max()) / float(ref.pow(2).mean().sqrt()) + + +def _math_state() -> dict: + return { + "float32_matmul_precision": torch.get_float32_matmul_precision(), + "cuda_matmul_allow_tf32": bool(torch.backends.cuda.matmul.allow_tf32), + "cudnn_allow_tf32": bool(torch.backends.cudnn.allow_tf32), + } + + +@pytest.fixture(scope="module", autouse=True) +def ieee_fp32(): + """The source's fp32 steps are IEEE fp32: refuse a process that substitutes TF32, both + before the module's tests and after them.""" + for var in ("TORCH_ALLOW_TF32_CUBLAS_OVERRIDE", "NVIDIA_TF32_OVERRIDE"): + assert var not in os.environ, ( + f"{var} is set and substitutes TF32 for fp32 GEMMs; run with `env -u {var}`" + ) + saved = ( + torch.get_float32_matmul_precision(), + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + ) + torch.set_float32_matmul_precision("highest") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + before = _fp32_gemm_probe() + assert before <= 1e-5, f"fp32 GEMM is not IEEE fp32 (relative error {before:.3g})" + state = {"math": _math_state(), "fp32_gemm_probe_before": before} + yield state + after = _fp32_gemm_probe() + state.update(math_after=_math_state(), fp32_gemm_probe_after=after) + _write_evidence("fp32-session", {"fixture": "ieee_fp32", "math": state}) + precision, allow_matmul, allow_cudnn = saved + torch.set_float32_matmul_precision(precision) + torch.backends.cuda.matmul.allow_tf32 = allow_matmul + torch.backends.cudnn.allow_tf32 = allow_cudnn + assert after <= 1e-5, f"fp32 GEMM stopped being IEEE fp32 (relative error {after:.3g})" + + +@pytest.fixture(scope="module", autouse=True) +def device_memory(): + """Hand the module's cached state and block back to the device when its tests are done.""" + yield + release_cache() + + +def _git_revision() -> str | None: + root = Path(tensorrt_llm.__file__).resolve().parents[1] + try: + out = subprocess.run( + ["git", "-C", str(root), "rev-parse", "HEAD"], + capture_output=True, + text=True, + timeout=30, + ) + except (OSError, subprocess.SubprocessError): + return None + return out.stdout.strip() or None + + +def _environment(ieee: dict) -> dict: + device = torch.cuda.current_device() + return { + "device_name": torch.cuda.get_device_name(device), + "device_capability": list(torch.cuda.get_device_capability(device)), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "tensorrt_llm": tensorrt_llm.__version__, + "tensorrt_llm_file": tensorrt_llm.__file__, + "source_revision": _git_revision(), + "python": sys.version.split()[0], + "machine": platform.machine(), + "argv": sys.argv, + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "slurm_step_id": os.environ.get("SLURM_STEP_ID"), + "execution": dict(EXECUTION), + "kv_cache": dict(KV_CACHE, layer_mask=[i for i, kept in enumerate(LAYER_MASK) if kept]), + "math": ieee, + } + + +def _write_evidence(name: str, record: dict) -> None: + directory = os.environ.get("BEAM_SPARSE_BLOCK_EVIDENCE") + if not directory: + return + path = Path(directory) + path.mkdir(parents=True, exist_ok=True) + (path / f"{name}.json").write_text(json.dumps(record, indent=1, default=str) + "\n") + + +def tensor_sha256(t: torch.Tensor) -> str: + flat = t.detach().contiguous().cpu().reshape(-1) + h = hashlib.sha256() + h.update(f"{t.dtype}|{tuple(t.shape)}|".encode()) + h.update(flat.view(torch.uint8).numpy().tobytes()) + return h.hexdigest() + + +# -------------------------------------------------- comparison + + +def compare(name: str, actual: torch.Tensor, expected: torch.Tensor, rule: str | None) -> dict: + """Score one bf16 tensor against one rule; `rule=None` measures without judging.""" + result = { + "tensor": name, + "rule": rule, + "shape": list(expected.shape), + "dtype": str(actual.dtype), + } + problems = [] + if tuple(actual.shape) != tuple(expected.shape): + problems.append(f"shape {tuple(actual.shape)} != {tuple(expected.shape)}") + if actual.dtype != DTYPE or expected.dtype != DTYPE: + problems.append(f"dtypes {actual.dtype}/{expected.dtype}, need {DTYPE}") + if problems: + result.update(passed=False, problems=problems) + return result + a, r = actual.double(), expected.to(actual.device).double() + if not bool(torch.isfinite(a).all()) or not bool(torch.isfinite(r).all()): + result.update(passed=False, problems=["non-finite values"]) + return result + diff = (a - r).abs() + numel = diff.numel() + result.update(numel=numel, unequal=int((a != r).sum()), max_abs=0.0) + if numel: + worst = int(diff.argmax()) + result.update( + max_abs=float(diff.max()), + mean_abs=float(diff.mean()), + worst_index=[int(i) for i in torch.unravel_index(torch.tensor(worst), diff.shape)], + worst_actual=float(a.reshape(-1)[worst]), + worst_expected=float(r.reshape(-1)[worst]), + ) + if rule is None: + result["passed"] = None + return result + atol, rtol = RULES[rule] + violations = int((diff > atol + rtol * r.abs()).sum()) + result.update(atol=atol, rtol=rtol, violations=violations) + if violations: + problems.append( + f"{violations} of {numel} elements outside atol {atol} + rtol {rtol} * |ref| " + f"(max abs {result['max_abs']:.3g})" + ) + result.update(passed=not problems, problems=problems) + return result + + +def compare_exact(name: str, actual: torch.Tensor, expected: torch.Tensor) -> dict: + """Integer selections: equal element for element and in order, never sorted first.""" + same_kind = tuple(actual.shape) == tuple(expected.shape) and actual.dtype == expected.dtype + if same_kind: + rows = (actual == expected).reshape(expected.shape[0], -1).all(dim=-1) + mismatched = [int(i) for i in torch.nonzero(~rows).flatten().cpu()] + else: + mismatched = list(range(expected.shape[0])) + problems = [] + if not same_kind: + problems.append( + f"{tuple(actual.shape)} {actual.dtype} != {tuple(expected.shape)} {expected.dtype}" + ) + elif mismatched: + problems.append(f"rows {mismatched[:8]} differ") + return { + "tensor": name, + "rule": "exact, in order", + "shape": list(expected.shape), + "mismatched_rows": mismatched, + "passed": not problems, + "problems": problems, + } + + +def inversion_gap(candidate: list, ref: list, scores: list, members: set) -> float: + """Largest ``scores[y] - scores[x]`` over the pairs of ``members`` that the candidate + ranks x before y and the reference ranks y before x; 0.0 when there is none. A rank is + the position in the ordered selection, and a member a selection lacks ranks after all of + it (8 for experts, 4 for groups).""" + rank_candidate = {m: i for i, m in enumerate(candidate)} + rank_ref = {m: i for i, m in enumerate(ref)} + last_candidate, last_ref = len(candidate), len(ref) + gap = 0.0 + for x in members: + for y in members: + before = rank_candidate.get(x, last_candidate) < rank_candidate.get(y, last_candidate) + after = rank_ref.get(y, last_ref) < rank_ref.get(x, last_ref) + if before and after: + gap = max(gap, scores[y] - scores[x]) + return gap + + +def judge_routing(own: dict, ref: dict) -> tuple[dict, torch.Tensor]: + """The component's routing against the whole reference block's, token by token. + + A token routed identically passes. Otherwise each decision taken differently has to be + a near tie in the reference's own scores: groups kept or ordered differently by their + ``group_scores``, experts selected or ordered differently by their ``selection`` scores + among the experts of the groups both kept (an expert of a group only one side kept is + the group decision's consequence). Returns the judgement and the mask of the tokens + routed to the same expert set. + """ + own_ids, ref_ids = own["ids"].tolist(), ref["ids"].tolist() + own_groups, ref_groups = own["groups"].tolist(), ref["groups"].tolist() + selection = ref["selection"].tolist() + group_scores = ref["group_scores"].tolist() + per_group = moe_reference.PER_GROUP + differing = [] + for t, (ids, ref_t) in enumerate(zip(own_ids, ref_ids)): + if ids == ref_t and own_groups[t] == ref_groups[t]: + continue + common = set(own_groups[t]) & set(ref_groups[t]) + members = {e for e in set(ids) | set(ref_t) if e // per_group in common} + differing.append( + { + "token": t, + "ids": ids, + "reference_ids": ref_t, + "groups": own_groups[t], + "reference_groups": ref_groups[t], + "group_gap": inversion_gap( + own_groups[t], ref_groups[t], group_scores[t], set(range(len(group_scores[t]))) + ), + "expert_gap": inversion_gap(ids, ref_t, selection[t], members), + } + ) + worst = max([0.0] + [max(d["group_gap"], d["expert_gap"]) for d in differing]) + beyond = [d["token"] for d in differing if max(d["group_gap"], d["expert_gap"]) > ROUTE_GAP] + matched = torch.tensor( + [set(ids) == set(ref_t) for ids, ref_t in zip(own_ids, ref_ids)], + dtype=torch.bool, + device=own["ids"].device, + ) + problems = [] + if beyond: + problems.append(f"tokens {beyond[:8]} route differently by up to {worst:.3g}") + judgement = { + "tensor": "routing_vs_reference_block", + "rule": f"every differing decision within {ROUTE_GAP:g} of the reference's own scores", + "tokens": len(ref_ids), + "differing_tokens": differing, + "set_matched_tokens": int(matched.sum()), + "max_gap": worst, + "passed": not problems, + "problems": problems, + } + return judgement, matched + + +def output_contract(out, tokens: int) -> dict: + problems = [] + if not isinstance(out, torch.Tensor): + problems.append(f"returned {type(out).__name__}") + elif tuple(out.shape) != (1, tokens, 5120) or out.dtype != DTYPE or not out.is_cuda: + problems.append(f"returned {tuple(out.shape)} {out.dtype} on {out.device}") + elif torch.cuda.get_device_capability(out.device) != SM: + problems.append(f"on {torch.cuda.get_device_name(out.device)}") + elif not bool(torch.isfinite(out).all()): + problems.append("non-finite values") + return { + "tensor": "output_contract", + "rule": f"[1, {tokens}, 5120] bf16 on the sm_103 device, finite", + "passed": not problems, + "problems": problems, + } + + +def tensor_versions(module: torch.nn.Module) -> dict[str, int]: + """The in-place write counter of every parameter and buffer, by name.""" + versions = {name: p._version for name, p in module.named_parameters()} + versions.update({name: b._version for name, b in module.named_buffers()}) + return versions + + +def unchanged(x: torch.Tensor, x_before: torch.Tensor, before: dict, after: dict) -> dict: + written = sorted(name for name, version in before.items() if after.get(name) != version) + problems = [] + if not bitwise_equal(x, x_before): + problems.append("the input was modified") + if written: + problems.append(f"{len(written)} parameters or buffers were written: {written[:4]}") + return { + "tensor": "input_and_parameters", + "rule": "input bits and parameter write counters unchanged by the call", + "passed": not problems, + "problems": problems, + } + + +def bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + """Same dtype, shape and bits (so -0.0 differs from +0.0).""" + if a.dtype != b.dtype or a.shape != b.shape: + return False + view = {torch.bfloat16: torch.int16, torch.float32: torch.int32}.get(a.dtype) + if view is None: + return bool(torch.equal(a, b)) + return bool(torch.equal(a.contiguous().view(view), b.contiguous().view(view))) + + +def assert_all_pass(results: list[dict]) -> None: + failed = [r for r in results if r["passed"] is False] + assert not failed, "\n".join( + f"{r['tensor']} [{r['rule']}]: {'; '.join(r['problems'])}" for r in failed + ) + + +# -------------------------------------------------- state, inputs, fresh context + +_STATE_CACHE: dict[int, dict] = {} + + +def release_cache() -> None: + """Drop the cached state and block and hand their memory back to the device.""" + _STATE_CACHE.clear() + gc.collect() + torch.cuda.empty_cache() + + +def generated(seed: int) -> dict: + """The seed's 782 tensors and [1, 129, 5120] input, generated once on the GPU.""" + if seed not in _STATE_CACHE: + release_cache() # one seed's 18.75 GiB of expert weights, and its block, at a time + state = reference.generate_state(seed) + problems = reference.state_problems(state, seed) + assert not problems, problems + _STATE_CACHE[seed] = { + "state": state, + "moe_state": reference.moe_state(state), + "input": reference.generate_input(seed), + "block": None, + "loaded": False, + } + return _STATE_CACHE[seed] + + +def loaded_block(seed: int) -> BeamSparseBlockBf16Sm103Tp1: + """The seed's cached block, loaded with the seed's state.""" + entry = generated(seed) + if entry["block"] is None: + entry["block"] = BeamSparseBlockBf16Sm103Tp1(block_config(), LAYER_IDX, device="cuda") + if not entry["loaded"]: + entry["block"].load_weights(entry["state"]) + entry["loaded"] = True + return entry["block"] + + +def cheap_weights() -> "OrderedDict[str, torch.Tensor]": + """A valid 782-tensor mapping that allocates nothing: one expanded constant per dtype, + so a load that copied any of it would change the generated values it replaced.""" + constants = { + torch.float32: torch.full((), 0.5, dtype=torch.float32, device="cuda"), + torch.bfloat16: torch.full((), 0.25, dtype=torch.bfloat16, device="cuda"), + } + return OrderedDict( + (name, constants[dtype].expand(shape)) + for name, (shape, dtype) in reference.param_shapes().items() + ) + + +@contextlib.contextmanager +def fresh_context(num_tokens: int, cached_tokens: int = 0, num_requests: int = 1, **mgr_overrides): + """A KVCacheManagerV2 whose layer mask selects logical layer 3 alone, and prepared + TrtllmAttentionMetadata for one new request.""" + mgr_args = dict( + num_layers=1, + layer_mask=list(LAYER_MASK), + num_kv_heads=BLOCK_CONFIG["num_key_value_heads"], + head_dim=BLOCK_CONFIG["head_dim"], + tokens_per_block=KV_CACHE["tokens_per_block"], + max_seq_len=KV_CACHE["max_seq_len"], + max_batch_size=num_requests, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.BF16, + ) + mgr_args.update(mgr_overrides) + mgr = KVCacheManagerV2( + KvCacheConfig( + max_tokens=KV_CACHE["num_blocks"] * KV_CACHE["tokens_per_block"] * num_requests, + enable_block_reuse=False, + ), + CacheTypeCpp.SELF, + **mgr_args, + ) + try: + request_ids = [REQUEST_ID + i for i in range(num_requests)] + mgr.add_dummy_requests(request_ids, token_nums=[cached_tokens + num_tokens] * num_requests) + md = TrtllmAttentionMetadata( + max_num_requests=num_requests, + max_num_tokens=block_module.MAX_TOKENS * num_requests, + kv_cache_manager=mgr, + num_contexts=num_requests, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=[cached_tokens] * num_requests + ), + seq_lens=torch.tensor([num_tokens] * num_requests, dtype=torch.int), + request_ids=request_ids, + prompt_lens=[cached_tokens + num_tokens] * num_requests, + ) + md.prepare() + yield md + finally: + mgr.shutdown() + + +def case_inputs(seed: int, length: int) -> tuple[torch.Tensor, torch.Tensor]: + x = generated(seed)["input"][:, :length].contiguous() + positions = torch.arange(length, device="cuda").view(1, length) + return x, positions + + +def reference_record(seed: int, x: torch.Tensor, positions: torch.Tensor) -> dict: + """The reference block of the seed's state on ``x``: its output and every named + intermediate, ``[T, ...]``, on the GPU.""" + record = {} + with torch.no_grad(): + out = reference.block(x, positions, generated(seed)["state"], record=record.__setitem__) + assert torch.equal(out[0], record["block_output"]) + return record + + +def run_block( + blk: BeamSparseBlockBf16Sm103Tp1, x: torch.Tensor, positions: torch.Tensor +) -> tuple[torch.Tensor, dict]: + """One forward on a fresh context: the output, the MoE's input (the pre-feed-forward + norm's output) and output from a forward hook, and the router op's mapping of the + call.""" + seen = [] + + def hook(module, args, output): + seen.append((args[0].clone(), output.clone())) + + handle = blk.mlp.register_forward_hook(hook) + try: + with fresh_context(x.shape[1]) as md, recording(beam_grouped_topk_router) as calls: + out = blk(x, positions, md) + torch.cuda.synchronize() + finally: + handle.remove() + assert len(seen) == 1 and len(calls) == 1, (len(seen), len(calls)) + ((pre_ff, moe_out),) = seen + return out, {"pre_ff": pre_ff, "moe_out": moe_out, "routing": calls[0]["output"]} + + +def block_results(seed: int, out: torch.Tensor, captured: dict, ref: dict) -> list[dict]: + """Every comparison of one forward against the reference block of the same input, + routing judged apart from numerical error.""" + moe_state = generated(seed)["moe_state"] + pre_ff, moe_out, routing = captured["pre_ff"], captured["moe_out"], captured["routing"] + tokens = pre_ff.shape[0] + results = [output_contract(out, tokens)] + # Same input, same routing: the reference routes the component's own MoE input exactly + # as the component did. + same_input = moe_reference.route( + pre_ff, moe_state["gate.weight"], moe_state["gate.e_score_correction_bias"] + ) + results.append(compare_exact("ids_on_own_moe_input", routing["ids"], same_input["ids"])) + results.append( + compare_exact("groups_on_own_moe_input", routing["groups"], same_input["groups"]) + ) + # Against the whole reference block: differing decisions are near ties. + judgement, matched = judge_routing(routing, ref["moe.routing"]) + results.append(judgement) + final = compare( + "block_output_on_set_matched_tokens", + out[0][matched], + ref["block_output"][matched], + "block_output", + ) + final["tokens"] = int(matched.sum()) + results.append(final) + results.append( + compare( + "moe_output_on_own_moe_input", + moe_out, + moe_reference.moe(pre_ff, moe_state), + "moe_output", + ) + ) + # Measured, not judged: how far the component's MoE input is from the reference's. + results.append( + compare("pre_feedforward_layernorm", pre_ff, ref["pre_feedforward_layernorm"], None) + ) + return results + + +# -------------------------------------------------- op-path instrumentation + + +@contextlib.contextmanager +def recording(*wrappers, experts: bool = False): + """Log every call of these entries, in order, with its nesting depth, the arguments that + identify the path and its output; with ``experts``, also every expert the MoE asks for + and the rows it passes.""" + calls: list[dict] = [] + depth = [0] + saved = [] + for w in wrappers: + cls = type(w) + original = cls.__call__ + saved.append((cls, "__call__", original)) + + def logged(self, *args, _original=original, _name=cls.__name__, **kwargs): + tensors = [a for a in args if isinstance(a, torch.Tensor)] + entry = { + "op": _name, + "depth": depth[0], + "shapes": [list(t.shape) for t in tensors], + } + entry.update({key: kwargs[key] for key in _RECORDED_KWARGS if key in kwargs}) + if "q" in kwargs: + entry["q_shape"] = list(kwargs["q"].shape) + entry["qkv"] = kwargs["q"] + calls.append(entry) + depth[0] += 1 + try: + out = _original(self, *args, **kwargs) + finally: + depth[0] -= 1 + entry["output"] = out + return out + + cls.__call__ = logged + if experts: + original_expert = BeamMoeBf16Sm103Tp1._expert + saved.append((BeamMoeBf16Sm103Tp1, "_expert", original_expert)) + + def asked(self, expert_id, rows, _original=original_expert): + calls.append({"op": "expert", "depth": depth[0], "expert_id": expert_id, "rows": rows}) + return _original(self, expert_id, rows) + + BeamMoeBf16Sm103Tp1._expert = asked + try: + yield calls + finally: + for cls, name, original in reversed(saved): + setattr(cls, name, original) + + +def loggable(calls: list[dict]) -> list[dict]: + """The recorded calls without their tensors, for the evidence record.""" + keep = ("op", "depth", "shapes", "q_shape", "expert_id", *_RECORDED_KWARGS) + out = [] + for call in calls: + entry = {key: call[key] for key in keep if key in call} + if "rows" in call: + entry["rows"] = list(call["rows"].shape) + out.append(entry) + return out + + +@contextlib.contextmanager +def builtin_beam_forbidden(): + """The built-in Beam modeling must not run while the component does.""" + targets = [ + (builtin_beam.BeamDecoderLayer, "forward"), + (builtin_beam.BeamAttention, "forward"), + (builtin_beam.BeamMoE, "forward"), + (builtin_beam.BeamMoE, "combine"), + (builtin_beam.BeamRouter, "forward"), + (builtin_beam.BeamRouter, "route"), + (builtin_beam.BeamSwiGLU, "forward"), + (builtin_beam.BeamRMSNorm, "forward"), + ] + saved = [(cls, name, getattr(cls, name)) for cls, name in targets] + + def refuse(self, *args, **kwargs): + raise AssertionError(f"the built-in {type(self).__name__} ran inside the component") + + for cls, name, _ in saved: + setattr(cls, name, refuse) + try: + yield + finally: + for cls, name, original in reversed(saved): + setattr(cls, name, original) + + +@contextlib.contextmanager +def no_dispatch(): + """Fail on any op call: rejected inputs must be rejected before a kernel runs.""" + saved = [] + for w in OPS: + cls = type(w) + saved.append((cls, cls.__call__)) + + def refuse(self, *args, _name=cls.__name__, **kwargs): + raise AssertionError(f"{_name} dispatched before the input was rejected") + + cls.__call__ = refuse + try: + yield + finally: + for cls, original in reversed(saved): + cls.__call__ = original + + +def cache_rows(md: TrtllmAttentionMetadata, length: int) -> tuple[torch.Tensor, torch.Tensor]: + """This request's K and V rows read back from the pool, [T, Hkv, D] each. The manager's + APIs take the logical layer, 3, and translate it to the local slot themselves.""" + mgr = md.kv_cache_manager + view = mgr.get_buffers(LAYER_IDX, kv_layout="HND") + view = view[:: mgr.get_layer_page_index_scale(LAYER_IDX) // 2] + tpb = KV_CACHE["tokens_per_block"] + keys, values = [], [] + for page in range((length + tpb - 1) // tpb): + k_off = int(md.kv_cache_block_offsets[0, 0, 0, page]) + rows = min(tpb, length - page * tpb) + keys.append(view[k_off // 2, 0, :, :rows, :].permute(1, 0, 2)) + values.append(view[k_off // 2, 1, :, :rows, :].permute(1, 0, 2)) + return torch.cat(keys), torch.cat(values) + + +# -------------------------------------------------- tests: configuration and construction + + +def test_configuration_is_the_fixed_scope() -> None: + """One configuration, no open choices: what the component accepts is what is tested.""" + assert CASES == ((17, 1), (17, 17), (17, 129), (41, 1), (41, 17), (41, 129)) + assert block_module.SUPPORTED_LAYER_IDX == LAYER_IDX + for field, value in block_module.SUPPORTED_CONFIG.items(): + assert BLOCK_CONFIG[field] == value, field + assert BLOCK_CONFIG["rope_parameters"]["sliding_attention"] == block_module.SUPPORTED_ROPE + assert BLOCK_CONFIG["layer_types"][LAYER_IDX] == "sliding_attention" + assert BLOCK_CONFIG["mlp_layer_types"][LAYER_IDX] == "sparse" + assert len(BLOCK_CONFIG["layer_types"]) == len(BLOCK_CONFIG["mlp_layer_types"]) == NUM_LAYERS + assert block_module.MAX_TOKENS == 1024 and max(LENGTHS) == reference.MAX_TOKENS + own = block_module.weight_shapes(block_config()) + assert list(own.items()) == list(reference.own_shapes().items()) + assert (reference.HEADS, reference.KV_HEADS, reference.HEAD_DIM) == (48, 8, 128) + assert reference.WINDOW == BLOCK_CONFIG["sliding_window"] == KV_CACHE["max_seq_len"] + assert reference.LAYER_IDX == LAYER_IDX and sum(LAYER_MASK) == 1 and LAYER_MASK[LAYER_IDX] + assert (reference.SEEDS, reference.LENGTHS) == (SEEDS, LENGTHS) + assert torch.cuda.get_device_capability() == SM, "this suite certifies sm_103 only" + + +def test_block_constructs_alone() -> None: + """The block is built directly from its configuration, without the model around it: its + twelve attention and norm tensors and the MoE component, no dense MLP, and block 3's + residual coefficients.""" + + def refuse(self, *args, **kwargs): + raise AssertionError(f"{type(self).__name__} constructed") + + with pytest.MonkeyPatch.context() as mp: + for cls in ( + builtin_beam.BeamForCausalLM, + builtin_beam.BeamModel, + builtin_beam.BeamDecoderLayer, + builtin_beam.BeamAttention, + builtin_beam.BeamMoE, + builtin_beam.BeamRouter, + builtin_beam.BeamSwiGLU, + ): + mp.setattr(cls, "__init__", refuse) + blk = BeamSparseBlockBf16Sm103Tp1(block_config(), LAYER_IDX, device="cuda") + assert isinstance(blk.mlp, BeamMoeBf16Sm103Tp1) + assert [name for name, _ in blk.named_children()] == ["mlp"] + tensors = blk.checkpoint_tensors() + names = list(tensors) + assert len(names) == 782 + assert names[:12] == list(reference.own_shapes()) + assert names[12:] == [reference.MLP_PREFIX + n for n in moe_reference.param_shapes()] + assert list(tensors) == list(reference.param_shapes()) + for dense in ("mlp.gate_proj.weight", "mlp.up_proj.weight", "mlp.down_proj.weight"): + assert dense not in tensors, dense + for name, (shape, dtype) in reference.param_shapes().items(): + t = tensors[name] + assert tuple(t.shape) == shape and t.dtype == dtype and t.is_cuda, name + assert tensors["mlp.gate.weight"] is blk.mlp.gate.weight + assert tensors["mlp.experts.255.down_proj.weight"] is blk.mlp.experts[255].down_proj + assert blk.attention_scales == (1 / math.sqrt(8), math.sqrt(7) / math.sqrt(8)) + assert blk.feed_forward_scales == (1 / 3, math.sqrt(8) / 3) + assert blk.attention_scales != block_module.depth_scales(0, True) + assert blk.feed_forward_scales != block_module.depth_scales(0, False) + assert "modeling_beam" not in Path(block_module.__file__).read_text() + + +# -------------------------------------------------- tests: declared limits + +_SLIDING_ROPE_1E6 = { + "sliding_attention": {"rope_type": "default", "rope_theta": 1e6}, + "full_attention": None, +} +_BAD_CONFIGS = [ + ( + "dense_block_3", + dict(mlp_layer_types=replaced(BLOCK_CONFIG["mlp_layer_types"], LAYER_IDX, "dense")), + "sparse", + ), + ( + "full_attention_block_3", + dict(layer_types=replaced(BLOCK_CONFIG["layer_types"], LAYER_IDX, "full_attention")), + "sliding_attention", + ), + ("rope_theta", dict(rope_parameters=_SLIDING_ROPE_1E6), "rope parameters"), + ("sliding_window", dict(sliding_window=1024), "unsupported sliding_window"), + ("head_dim", dict(head_dim=64), "unsupported head_dim"), +] + + +@pytest.mark.parametrize("case,overrides,match", _BAD_CONFIGS, ids=[c[0] for c in _BAD_CONFIGS]) +def test_rejects_unsupported_configuration(case, overrides, match) -> None: + """Anything outside the one configuration is refused at construction.""" + with pytest.raises(ValueError, match=match): + BeamSparseBlockBf16Sm103Tp1(block_config(**overrides), LAYER_IDX, device="cuda") + + +def test_rejects_other_blocks_and_devices() -> None: + for layer_idx in (0, 2, 6): + with pytest.raises(ValueError, match="only sparse blocks"): + BeamSparseBlockBf16Sm103Tp1(block_config(), layer_idx, device="cuda") + with pytest.raises(ValueError, match="sm_103"): + BeamSparseBlockBf16Sm103Tp1(block_config(), LAYER_IDX, device="cpu") + + +# -------------------------------------------------- tests: loading + + +def test_load_weights_maps_every_tensor() -> None: + """Exactly the 782 names, each copied exactly into its own tensor: the twelve attention + and norm tensors into the block, the ``mlp.`` ones into the MoE's router and into expert + e for every e.""" + seed = 17 + state = generated(seed)["state"] + blk = loaded_block(seed) + loaded = blk.checkpoint_tensors() + assert list(loaded) == list(state) == list(reference.param_shapes()) + wrong = [] + for name, (shape, dtype) in reference.param_shapes().items(): + t = loaded[name] + if tuple(t.shape) != shape or t.dtype != dtype or not t.is_cuda: + wrong.append(f"{name}: {tuple(t.shape)} {t.dtype} on {t.device}") + elif not torch.equal(t, state[name]): + wrong.append(f"{name}: values differ") + assert t.data_ptr() != state[name].data_ptr(), f"{name} aliases the given tensor" + assert not wrong, wrong[:8] + moe = blk.mlp + for expert_id, expert in enumerate(moe.experts): + for attr in ("gate_proj", "up_proj", "down_proj"): + given = state[f"mlp.experts.{expert_id}.{attr}.weight"] + assert torch.equal(getattr(expert, attr), given), (expert_id, attr) + assert moe.gate.weight.dtype == moe.gate.e_score_correction_bias.dtype == torch.float32 + assert torch.equal(moe.gate.weight, state["mlp.gate.weight"]) + assert torch.equal(moe.gate.e_score_correction_bias, state["mlp.gate.e_score_correction_bias"]) + + +_BAD_LOADS = [ + ("late_expert_shape", ValueError, r"experts\.255\.up_proj\.weight must be \[2560, 5120\]"), + ("q_norm_fp32", TypeError, r"self_attn\.q_norm\.weight must be torch\.bfloat16"), + ("router_weight_bf16", TypeError, r"gate\.weight must be torch\.float32"), +] + + +@pytest.mark.parametrize("case,error,match", _BAD_LOADS, ids=[c for c, _, _ in _BAD_LOADS]) +def test_rejected_load_changes_nothing(case, error, match) -> None: + """Every one of the 782 tensors is checked before any is copied: a mapping malformed + late is rejected without writing any parameter or buffer, every value stays, and the + block still runs.""" + seed = 17 + entry = generated(seed) + blk = loaded_block(seed) + bad = cheap_weights() + if case == "late_expert_shape": + bad["mlp.experts.255.up_proj.weight"] = bad["mlp.experts.255.up_proj.weight"].t() + elif case == "q_norm_fp32": + bad["self_attn.q_norm.weight"] = torch.zeros((), device="cuda").expand(128) + elif case == "router_weight_bf16": + bad["mlp.gate.weight"] = bad["mlp.gate.weight"].to(DTYPE) + names = ( + "self_attn.q_norm.weight", + "post_feedforward_layernorm.weight", + "mlp.gate.weight", + "mlp.gate.e_score_correction_bias", + "mlp.experts.255.up_proj.weight", + ) + snapshots = {name: blk.checkpoint_tensors()[name].clone() for name in names} + versions = tensor_versions(blk) + try: + with no_dispatch(), pytest.raises(error, match=match): + blk.load_weights(bad) + finally: + written = sorted(n for n, v in tensor_versions(blk).items() if v != versions[n]) + if written: + entry["loaded"] = False # reload the generated state before the block runs again + current = blk.checkpoint_tensors() + changed = [name for name, value in snapshots.items() if not torch.equal(current[name], value)] + assert not written, f"{len(written)} tensors written by a rejected load: {written[:8]}" + assert not changed, changed + x, positions = case_inputs(seed, 1) + with fresh_context(1) as md: + out = blk(x, positions, md) + assert out.shape == (1, 1, 5120) and bool(torch.isfinite(out).all()) + + +_BAD_INPUTS = [ + "fp16_hidden", + "fp32_hidden", + "hidden_5120x2", + "batch_2", + "too_long", + "empty", + "non_contiguous", + "positions_offset", + "float_positions", + "positions_shape", + "metadata_type", + "cached_prefix", + "two_requests", + "context_length_mismatch", + "kv_cache_heads", + "kv_cache_dtype", + "layer0_cache", +] + + +@pytest.mark.parametrize("case", _BAD_INPUTS) +def test_rejects_unsupported_inputs_before_dispatch(case) -> None: + """Each declared restriction raises ValueError on the loaded block, and no op runs + first -- including a cache manager that does not hold logical layer 3 (a default + one-layer manager holds layer 0).""" + blk = loaded_block(17) + length = 17 + x = torch.randn(1, length, 5120, device="cuda").to(DTYPE) + positions = torch.arange(length, device="cuda").view(1, length) + context = dict(num_tokens=length) + if case == "fp16_hidden": + x = x.half() + elif case == "fp32_hidden": + x = x.float() + elif case == "hidden_5120x2": + x = torch.cat((x, x), dim=-1) + elif case == "batch_2": + x = torch.cat((x, x), dim=0) + elif case == "too_long": + length = block_module.MAX_TOKENS + 1 + x = torch.randn(1, length, 5120, device="cuda").to(DTYPE) + positions = torch.arange(length, device="cuda").view(1, length) + context = dict(num_tokens=17) + elif case == "empty": + x = x[:, :0] + positions = positions[:, :0] + elif case == "non_contiguous": + x = torch.randn(1, length, 2 * 5120, device="cuda").to(DTYPE)[..., ::2] + elif case == "positions_offset": + positions = positions + 1 + elif case == "float_positions": + positions = positions.float() + elif case == "positions_shape": + positions = positions.view(length) + elif case == "cached_prefix": + context = dict(num_tokens=length, cached_tokens=8) + elif case == "two_requests": + context = dict(num_tokens=length, num_requests=2) + elif case == "context_length_mismatch": + context = dict(num_tokens=length - 1) + elif case == "kv_cache_heads": + context = dict(num_tokens=length, num_kv_heads=4) + elif case == "kv_cache_dtype": + context = dict(num_tokens=length, dtype=DataType.HALF) + elif case == "layer0_cache": + context = dict(num_tokens=length, layer_mask=None) + with fresh_context(**context) as md: + if case == "metadata_type": + md = object() + elif case == "layer0_cache": + assert dict(md.kv_cache_manager.layer_offsets) == {0: 0} + with no_dispatch(), pytest.raises(ValueError): + blk(x, positions, md) + + +def test_cache_slot_is_the_managers_local_index(ieee_fp32) -> None: + """Logical layer 3 lives in slot 0 of a manager whose layer mask selects it alone: the + attention op is called with local index 0, and the K and V it appended are read back + through the manager under the logical index 3.""" + seed, length = 17, 17 + blk = loaded_block(seed) + x, positions = case_inputs(seed, length) + with fresh_context(length) as md: + mgr = md.kv_cache_manager + layout = { + "manager": type(mgr).__name__, + "num_layers": mgr.num_layers, + "pp_layers": list(mgr.pp_layers), + "layer_offsets": dict(mgr.layer_offsets), + "num_local_layers": mgr.num_local_layers, + } + step = blk.step_args(md) + with recording(thop_attention) as calls: + out = blk(x, positions, md) + torch.cuda.synchronize() + keys, values = cache_rows(md, length) + _write_evidence( + f"cache-slot-seed{seed}-T{length}", + { + "test": "test_cache_slot_is_the_managers_local_index", + "layout": layout, + "step_local_layer_idx": step["local_layer_idx"], + "op_calls": loggable(calls), + "environment": _environment(ieee_fp32), + }, + ) + assert isinstance(mgr, KVCacheManagerV2) + assert layout["num_layers"] == NUM_LAYERS and layout["pp_layers"] == [LAYER_IDX] + assert layout["layer_offsets"] == {LAYER_IDX: 0} and layout["num_local_layers"] == 1 + assert step["local_layer_idx"] == 0 + (call,) = calls + assert { + k: call[k] + for k in ( + "local_layer_idx", + "num_heads", + "num_kv_heads", + "head_size", + "attention_window_size", + "mask_type", + "is_fused_qkv", + ) + } == { + "local_layer_idx": 0, + "num_heads": 48, + "num_kv_heads": 8, + "head_size": 128, + "attention_window_size": 512, + "mask_type": 1, + "is_fused_qkv": True, + } + assert call["q_shape"] == [length, (48 + 2 * 8) * 128] + assert out.shape == (1, length, 5120) and bool(torch.isfinite(out).all()) + # The op appended exactly this call's post-RoPE K and V to logical layer 3's slot. + qkv = call["qkv"] + k_in = qkv[:, 48 * 128 : 56 * 128].reshape(length, 8, 128) + v_in = qkv[:, 56 * 128 :].reshape(length, 8, 128) + assert torch.equal(keys, k_in) and torch.equal(values, v_in) + + +def test_cache_slot_in_a_manager_holding_two_layers(ieee_fp32) -> None: + """In a manager that holds logical layers 2 and 3, layer 3 is local slot 1: the attention + op is called with that slot, the K and V are appended there, and the block output is + bitwise the output on a manager holding layer 3 alone.""" + seed, length = 17, 17 + blk = loaded_block(seed) + x, positions = case_inputs(seed, length) + with fresh_context(length) as md: + alone = blk(x, positions, md) + with fresh_context( + length, num_layers=2, layer_mask=[i in (2, 3) for i in range(NUM_LAYERS)] + ) as md: + assert dict(md.kv_cache_manager.layer_offsets) == {2: 0, 3: 1} + assert blk.step_args(md)["local_layer_idx"] == 1 + with recording(thop_attention) as calls: + out = blk(x, positions, md) + torch.cuda.synchronize() + keys, values = cache_rows(md, length) + (call,) = calls + assert call["local_layer_idx"] == 1 + qkv = call["qkv"] + assert torch.equal(keys, qkv[:, 48 * 128 : 56 * 128].reshape(length, 8, 128)) + assert torch.equal(values, qkv[:, 56 * 128 :].reshape(length, 8, 128)) + assert bitwise_equal(out, alone) + + +# -------------------------------------------------- tests: the six cases + + +@pytest.mark.parametrize("seed,length", CASES, ids=CASE_IDS) +def test_block_matches_reference(seed, length, ieee_fp32) -> None: + """One case, routing judged apart from numerical error: exact routing on the block's own + MoE input; against the whole reference block, routing differences only on near ties and + the output within its rule on the tokens routed alike; the MoE output within its rule + on the block's own MoE input; input and parameters untouched.""" + blk = loaded_block(seed) + x, positions = case_inputs(seed, length) + x_before = x.clone() + versions = tensor_versions(blk) + out, captured = run_block(blk, x, positions) + after = tensor_versions(blk) + ref = reference_record(seed, x, positions) + results = block_results(seed, out, captured, ref) + results.append(unchanged(x, x_before, versions, after)) + _write_evidence( + f"block-seed{seed}-T{length}", + { + "test": "test_block_matches_reference", + "case": {"seed": seed, "length": length}, + "rules": {k: list(v) for k, v in RULES.items()}, + "route_gap": ROUTE_GAP, + "results": results, + "input_sha256": tensor_sha256(x), + "environment": _environment(ieee_fp32), + }, + ) + assert_all_pass(results) + + +# -------------------------------------------------- tests: the component path + + +def test_forward_runs_the_component_path(ieee_fp32) -> None: + """One case: the component's own composition ran, op by op, with the experts' calls + nested inside BeamDispatchCombine, through thop attention on this GPU, with no SDPA, no + built-in Beam module and every op's guard admitting its arguments.""" + seed, length = 17, 17 + blk = loaded_block(seed) + x, positions = case_inputs(seed, length) + activities = [torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] + with fresh_context(length) as md: + with builtin_beam_forbidden(), validating(*OPS), recording(*OPS, experts=True) as calls: + with torch.profiler.profile(activities=activities) as prof: + out = blk(x, positions, md) + torch.cuda.synchronize() + manager = md.kv_cache_manager + events = prof.events() + kernels = [e.name for e in events if e.device_type == torch.autograd.DeviceType.CUDA] + cpu_ops = sorted({e.name for e in events if e.device_type == torch.autograd.DeviceType.CPU}) + names = [c["op"] for c in calls] + record = { + "test": "test_forward_runs_the_component_path", + "case": {"seed": seed, "length": length}, + "block": f"{type(blk).__module__}.{type(blk).__qualname__}", + "op_calls": loggable(calls), + "cuda_kernels": kernels, + "attention_kernels": [k for k in kernels if _ATTENTION_KERNEL.search(k)], + "sdpa_cpu_ops": [n for n in cpu_ops if _SDPA_OPS.search(n)], + "kv_cache_manager": type(manager).__name__, + "environment": _environment(ieee_fp32), + } + _write_evidence(f"path-seed{seed}-T{length}", record) + + assert type(blk).__module__.endswith("sparse_block_bf16__sm_103__tp1.block") + top = [c for c in calls if c["depth"] == 0] + assert tuple(c["op"] for c in top) == EXPECTED_CALLS, [c["op"] for c in top] + # Every other call ran inside BeamDispatchCombine: right after it, one level down. + position = names.index("BeamDispatchCombine") + inner = len(calls) - len(top) + depths = [c["depth"] for c in calls] + assert depths == [0] * (position + 1) + [1] * inner + [0] * (len(calls) - position - 1 - inner) + router = calls[names.index("BeamGroupedTopkRouter")] + ids = router["output"]["ids"] + selected = sorted(set(ids.reshape(-1).tolist())) + nested = calls[position + 1 : position + 1 + inner] + assert [c["op"] for c in nested] == list(EXPERT_CALLS) * len(selected) + assert [c["expert_id"] for c in nested if c["op"] == "expert"] == selected + assert sum(c["rows"].shape[0] for c in nested if c["op"] == "expert") == 8 * length + assert inner == len(EXPERT_CALLS) * len(selected) + (call,) = [c for c in calls if c["op"] == "ThopAttention"] + assert call["local_layer_idx"] == 0 and call["num_contexts"] == 1 + assert call["num_ctx_tokens"] == length and call["use_paged_context_fmha"] is False + assert call["quant_mode"] == 0 and call["position_embedding_type"] == 0 + assert isinstance(manager, KVCacheManagerV2) + assert record["attention_kernels"], "no attention kernel ran on the GPU" + assert not record["sdpa_cpu_ops"], record["sdpa_cpu_ops"] + assert torch.cuda.get_device_capability(x.device) == SM + assert out.shape == (1, length, 5120) and bool(torch.isfinite(out.float()).all()) + + +# -------------------------------------------------- tests: fixtures + + +def test_causal_prefix_independence(ieee_fp32) -> None: + """Two 129-token inputs sharing their first 17 rows: those rows reach the MoE bit for + bit alike and are routed alike, and their outputs agree within the block rule (each + expert's batch differs, so only the expert GEMMs' row counts change).""" + seed, length, prefix = 17, 129, 17 + blk = loaded_block(seed) + x_a, positions = case_inputs(seed, length) + suffix = reference.generate_input(41)[:, prefix:length] + x_b = torch.cat((x_a[:, :prefix], suffix), dim=1).contiguous() + assert not torch.equal(x_b[:, prefix:], x_a[:, prefix:]) + out_a, captured_a = run_block(blk, x_a, positions) + out_b, captured_b = run_block(blk, x_b, positions) + ids_a, ids_b = captured_a["routing"]["ids"], captured_b["routing"]["ids"] + results = [ + compare_exact("prefix_ids", ids_b[:prefix], ids_a[:prefix]), + { + "tensor": "prefix_moe_input", + "rule": "bitwise", + "passed": bitwise_equal(captured_b["pre_ff"][:prefix], captured_a["pre_ff"][:prefix]), + "problems": [], + }, + compare("prefix_block_output", out_b[0, :prefix], out_a[0, :prefix], "block_output"), + ] + if not results[1]["passed"]: + results[1]["problems"].append("the shared rows' MoE inputs differ") + _write_evidence( + f"prefix-seed{seed}-T{length}", + { + "test": "test_causal_prefix_independence", + "prefix": prefix, + "results": results, + "suffix_ids_differ": not torch.equal(ids_b[prefix:], ids_a[prefix:]), + "environment": _environment(ieee_fp32), + }, + ) + assert_all_pass(results) + assert not torch.equal(out_b[0, prefix:], out_a[0, prefix:]), "the suffix changed nothing" + + +def test_residual_uses_block3_coefficients(ieee_fp32, monkeypatch) -> None: + """The residual entry with the block's coefficients is block 3's add bit for bit and + differs from block 0's; and the block comparison rejects a block run with block 0's + coefficients, and accepts the unchanged and the restored block.""" + seed, length = 17, 17 + blk = loaded_block(seed) + gen = torch.Generator(device="cuda") + gen.manual_seed(moe_reference.tensor_seed(seed, "residual.branch")) + branch = torch.randn(length, 5120, generator=gen, device="cuda").to(DTYPE) + gen.manual_seed(moe_reference.tensor_seed(seed, "residual.stream")) + stream = torch.randn(length, 5120, generator=gen, device="cuda").to(DTYPE) + assert blk.attention_scales == reference.ATTENTION_SCALES + assert blk.feed_forward_scales == reference.FEED_FORWARD_SCALES + for label, scales, is_attention in ( + ("attention", blk.attention_scales, True), + ("feed_forward", blk.feed_forward_scales, False), + ): + got = beam_depth_residual(branch, stream, *scales) + expected = reference.residual_add(branch, stream, reference.depth_coefs(3, is_attention)) + block0 = reference.residual_add(branch, stream, reference.depth_coefs(0, is_attention)) + assert bitwise_equal(got, expected), label + assert not torch.equal(got, block0), label + + x, positions = case_inputs(seed, length) + ref = reference_record(seed, x, positions) + runs = {} + for label in ("unchanged", "block0_coefficients", "restored"): + if label == "block0_coefficients": + monkeypatch.setattr(blk, "attention_scales", block_module.depth_scales(0, True)) + monkeypatch.setattr(blk, "feed_forward_scales", block_module.depth_scales(0, False)) + elif label == "restored": + monkeypatch.undo() + out, captured = run_block(blk, x, positions) + runs[label] = block_results(seed, out, captured, ref) + _write_evidence( + f"fault-residual-seed{seed}-T{length}", + { + "test": "test_residual_uses_block3_coefficients", + "fault": "block 0's residual coefficients in both adds, test-scoped", + "runs": runs, + "environment": _environment(ieee_fp32), + }, + ) + assert_all_pass(runs["unchanged"]) + failed = {r["tensor"] for r in runs["block0_coefficients"] if r["passed"] is False} + assert failed, "block 0's coefficients passed the block comparison" + assert failed <= {"routing_vs_reference_block", "block_output_on_set_matched_tokens"}, failed + assert blk.attention_scales == block_module.depth_scales(LAYER_IDX, True) + assert blk.feed_forward_scales == block_module.depth_scales(LAYER_IDX, False) + assert_all_pass(runs["restored"]) + + +def test_changed_routing_between_block_calls(ieee_fp32) -> None: + """Input A, then input B, then A again, each on a fresh context: B routes differently + from A and meets the block rules itself, and the repeated A reproduces the first one bit + for bit -- its output, MoE input and MoE output.""" + seed, length = 17, 17 + blk = loaded_block(seed) + x_a, positions = case_inputs(seed, length) + x_b = reference.generate_input(41)[:, :length].contiguous() + runs = [run_block(blk, x, positions) for x in (x_a, x_b, x_a)] + (out_a1, cap_a1), (out_b, cap_b), (out_a2, cap_a2) = runs + results_b = block_results(seed, out_b, cap_b, reference_record(seed, x_b, positions)) + ids_a1, ids_b, ids_a2 = [c["routing"]["ids"] for c in (cap_a1, cap_b, cap_a2)] + differing = [int(t) for t in torch.nonzero((ids_b != ids_a1).any(dim=-1)).flatten().cpu()] + _write_evidence( + f"changed-routing-seed{seed}-T{length}", + { + "test": "test_changed_routing_between_block_calls", + "tokens_routed_differently_in_b": differing, + "results_b": results_b, + "environment": _environment(ieee_fp32), + }, + ) + assert differing, "input B routed every token exactly as input A" + assert_all_pass(results_b) + assert torch.equal(ids_a2, ids_a1) + assert bitwise_equal(out_a2, out_a1) + assert bitwise_equal(cap_a2["pre_ff"], cap_a1["pre_ff"]) + assert bitwise_equal(cap_a2["moe_out"], cap_a1["moe_out"]) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/test_loading_contracts.py b/tests/unittest/_torch/modeling_v2/models/beam/test_loading_contracts.py new file mode 100644 index 000000000000..871f3897e6c6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/test_loading_contracts.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU fault-injection checks for weight loading; no checkpoint weights or collectives.""" + +import json +from collections import OrderedDict +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_bf16__sm_103__tp4 import ( # noqa: E501 + model as bf16, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.checkpoint_layout_nvfp4__sm_103__tp4 import ( # noqa: E501 + model as nvfp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1 import ( # noqa: E501 + block as dense, +) +from tensorrt_llm._torch.models import modeling_beam + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture(params=[bf16, nvfp4], ids=["bf16", "nvfp4"]) +def layout_loader(request, monkeypatch, tmp_path): + module = request.param + cls = ( + bf16.BeamCheckpointLayoutBf16Sm103Tp4 + if module is bf16 + else nvfp4.BeamCheckpointLayoutNvfp4Sm103Tp4 + ) + (tmp_path / "config.json").write_text("{}") + (tmp_path / "hf_quant_config.json").write_text("{}") + spec = SimpleNamespace( + name="weight", bank=None, expert=None, interval=lambda rank: (0, 2), local_shape=(2,) + ) + dst = torch.zeros(2, dtype=torch.bfloat16) + component = SimpleNamespace( + config={}, + rank=0, + layers=(), + expert_names=(), + specs=OrderedDict(weight=spec), + usable=True, + load_record={"old": True}, + tensor=lambda name: dst, + ) + component.schedule = SimpleNamespace(is_sparse=lambda layer: False) + agreements = [] + + def agree(ok): + agreements.append(ok) + return ok + + component._agree = agree + catalog = SimpleNamespace( + entries={"weight": SimpleNamespace(path="part", begin=0)}, + validate=lambda *args, **kwargs: None, + ) + monkeypatch.setattr(module, "SafetensorsCatalog", lambda path: catalog) + monkeypatch.setattr(module, "tensor_specs", lambda config: {}) + if module is nvfp4: + monkeypatch.setattr(module, "check_recipe", lambda *args: {}) + monkeypatch.setattr(module, "stored_entries", lambda *args: {}) + monkeypatch.setattr(module, "validate", lambda *args, **kwargs: None) + reader = SimpleNamespace(bytes_read=4, reads=1, closed=False) + + def close(): + reader.closed = True + + reader.close = close + monkeypatch.setattr(module, "TileReader", lambda: reader) + monkeypatch.setattr(module, "read_slice_into", lambda *args, **kwargs: dst.fill_(1)) + return module, cls.load, component, agreements, catalog, reader, tmp_path + + +@pytest.mark.parametrize("stored", [None, ["bad"]]) +def test_malformed_config_reaches_validation_agreement(layout_loader, stored): + _, load, component, agreements, _, _, path = layout_loader + (path / "config.json").write_text(json.dumps(stored)) + with pytest.raises(bf16.LoadError) as exc: + load(component, str(path)) + assert isinstance(exc.value.__cause__, (TypeError, AttributeError)) + assert agreements == [False] + assert component.usable # Validation rejected the reload before touching storage. + assert component.load_record == {"old": True} + + +@pytest.mark.parametrize("stage", ["catalog", "reader", "ordering", "read", "close"]) +def test_rank_local_exception_reaches_agreement(layout_loader, monkeypatch, stage): + module, load, component, agreements, catalog, reader, path = layout_loader + failure = TypeError("injected rank-local failure") + + def fail(*args, **kwargs): + raise failure + + if stage == "catalog": + monkeypatch.setattr(module, "SafetensorsCatalog", fail) + elif stage == "reader": + monkeypatch.setattr(module, "TileReader", fail) + elif stage == "ordering": + catalog.entries.clear() + elif stage == "read": + monkeypatch.setattr(module, "read_slice_into", fail) + else: + reader.close = fail + with pytest.raises(bf16.LoadError) as exc: + load(component, str(path)) + assert agreements == ([False] if stage == "catalog" else [True, False]) + if stage == "ordering": + assert isinstance(exc.value.__cause__, KeyError) + else: + assert exc.value.__cause__ is failure + if stage != "catalog": + assert not component.usable and component.load_record is None + if stage in ("ordering", "read"): + assert reader.closed + + +def test_successful_load_and_peer_failure(layout_loader): + _, load, component, agreements, _, reader, path = layout_loader + record = load(component, str(path)) + assert agreements == [True, True] + assert component.usable and record["tensors"] == 1 + assert reader.closed and torch.equal(component.tensor("weight"), torch.ones(2)) + agreements.clear() + + def peer_fails_completion(ok): + agreements.append(ok) + return len(agreements) == 1 + + component._agree = peer_fails_completion + with pytest.raises(bf16.LoadError, match="every rank"): + load(component, str(path)) + assert agreements == [True, True] + assert not component.usable and component.load_record is None + + +@pytest.fixture +def tiny_dense(monkeypatch): + # Exercise the loader with tiny CPU storage, without allocating a GB300 block. + block = dense.BeamDenseBlockBf16Sm103Tp1.__new__(dense.BeamDenseBlockBf16Sm103Tp1) + torch.nn.Module.__init__(block) + block._shapes = OrderedDict([("input_layernorm.weight", (2,)), ("mlp.up_proj.weight", (2, 2))]) + for name, shape in block._shapes.items(): + block.register_parameter( + dense._attribute(name), + torch.nn.Parameter(torch.zeros(shape, dtype=torch.bfloat16), requires_grad=False), + ) + block._loaded = False + block.layer_idx, block.num_kv_heads, block.head_dim = 0, 1, 2 + monkeypatch.setattr(dense._attention, "check_step", lambda *args, **kwargs: None) + return block + + +def _weights(block, value): + return { + name: torch.full(shape, value, dtype=torch.bfloat16) + for name, shape in block._shapes.items() + } + + +@pytest.mark.parametrize("bad_kind", ["shape", "dtype", "object", "meta"]) +def test_rejected_dense_reload_preserves_weights(tiny_dense, bad_kind): + block = tiny_dense + good = _weights(block, 1) + block.load_weights(good) + invalid = _weights(block, 2) + invalid["mlp.up_proj.weight"] = { + "shape": torch.zeros(3, dtype=torch.bfloat16), + "dtype": torch.zeros(2, 2, dtype=torch.float32), + "object": None, + "meta": torch.empty(2, 2, dtype=torch.bfloat16, device="meta"), + }[bad_kind] + with pytest.raises((TypeError, ValueError)): + block.load_weights(invalid) + assert block._loaded + assert all( + torch.equal(tensor, good[name]) for name, tensor in block.checkpoint_tensors().items() + ) + + +def test_dense_cannot_run_before_load_or_after_failed_copy(tiny_dense, monkeypatch): + block = tiny_dense + with pytest.raises(ValueError, match="load_weights"): + block(None, None, None) + block.load_weights(_weights(block, 1)) + original_copy = torch.Tensor.copy_ + + def fail_second(dst, src, *args, **kwargs): + if dst is block.mlp_up_proj: + raise RuntimeError("injected copy failure") + return original_copy(dst, src, *args, **kwargs) + + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "copy_", fail_second) + with pytest.raises(RuntimeError, match="copy failure"): + block.load_weights(_weights(block, 2)) + assert not block._loaded + with pytest.raises(ValueError, match="load_weights"): + block(None, None, None) + block.load_weights(_weights(block, 3)) + assert block._loaded + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_builtin_loader_rejects_tp_before_copy(monkeypatch, tp_size): + calls = [] + monkeypatch.setattr(modeling_beam, "load_beam_state_dict", lambda *args: calls.append(args)) + model = SimpleNamespace(model_config=SimpleNamespace(mapping=SimpleNamespace(tp_size=tp_size))) + if tp_size > 1: + with pytest.raises(ValueError, match="unsharded"): + modeling_beam.BeamForCausalLM.load_weights(model, {}) + assert calls == [] + else: + modeling_beam.BeamForCausalLM.load_weights(model, {}) + assert calls == [(model, {})] diff --git a/tests/unittest/_torch/modeling_v2/residual/test_modeling_v2_beam_depth_residual.py b/tests/unittest/_torch/modeling_v2/residual/test_modeling_v2_beam_depth_residual.py new file mode 100644 index 000000000000..670b42253440 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/residual/test_modeling_v2_beam_depth_residual.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_depth_residual entry.""" + +import math +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.residual.beam_depth_residual import ( + beam_depth_residual as op, +) + +__extra_import_path__ = [".."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_depth_residual requires a CUDA device" + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + shape = (spec["tokens"], spec["hidden"]) + branch = torch.randn(shape, generator=g, device="cuda").to(spec["dtype"]) + residual = (2.0 * torch.randn(shape, generator=g, device="cuda")).to(spec["dtype"]) + return branch, residual + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bitwise: two scalings and an add, each an fp32 operation rounded once.""" + certifying(op) + branch, residual = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + branch_scale, residual_scale = cell.spec["scales"] + out = op(branch, residual, branch_scale, residual_scale) + assert out.shape == branch.shape and out.dtype == branch.dtype + ref = op.reference(branch, residual, branch_scale, residual_scale) + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) + + +def test_scales_are_the_source_depth_coefficients() -> None: + """The cells carry blocks 0 and 3's coefficients: the n-th residual add scales by + (1, sqrt(n - 1)) / sqrt(n); block i's attention add is n = 2(i + 1), its feed-forward + add n = 2(i + 1) + 1.""" + scales = {cell.why.split(",")[0]: cell.spec["scales"] for cell in op.CELLS} + assert scales["block 0 attention add"] == (1 / math.sqrt(2), 1 / math.sqrt(2)) + assert scales["block 0 feed-forward add"] == (1 / math.sqrt(3), math.sqrt(2) / math.sqrt(3)) + assert scales["block 3 attention add"] == (1 / math.sqrt(8), math.sqrt(7) / math.sqrt(8)) + assert scales["block 3 feed-forward add"] == (1 / 3, math.sqrt(8) / 3) + for block, (attention, feed_forward) in ((0, (2, 3)), (3, (8, 9))): + for kind, n in (("attention", attention), ("feed-forward", feed_forward)): + branch, residual = scales[f"block {block} {kind} add"] + assert (branch, residual) == (1 / math.sqrt(n), math.sqrt(n - 1) / math.sqrt(n)) + + +def test_fused_fp32_add_is_a_different_answer() -> None: + branch, residual = _build(dict(tokens=129, hidden=5120, dtype=torch.bfloat16), seed=5) + a, b = 1 / math.sqrt(3), math.sqrt(2) / math.sqrt(3) + fused = (branch.float() * a + residual.float() * b).to(torch.bfloat16) + assert not torch.equal(op(branch, residual, a, b), fused) + + +@pytest.mark.parametrize( + "case,match", + [ + ("broadcast_residual", "must match"), + ("fp32_residual", "share one dtype"), + ("tensor_scale", "Python floats"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + branch, residual = _build(dict(tokens=17, hidden=5120, dtype=torch.bfloat16), seed=3) + branch_scale, residual_scale = 1 / math.sqrt(2), 1 / math.sqrt(2) + if case == "broadcast_residual": + residual = residual[:1] + elif case == "fp32_residual": + residual = residual.float() + else: + branch_scale = torch.tensor(branch_scale, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match=match): + with validating(op): + op(branch, residual, branch_scale, residual_scale)