From 4380dff757e5aa51f07e3c46ad3c7bd2183dbd0f Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Fri, 16 Oct 2026 16:10:41 -0700 Subject: [PATCH 01/13] [None][feat] Add initial Beam component implementation Signed-off-by: Maxim Gluhovskoi --- .../backends/fmha/flashinfer_trtllm_gen.py | 35 + tensorrt_llm/_torch/models/modeling_beam.py | 848 ++++++++++++++++++ .../test_flashinfer_trtllm_gen_fmha.py | 46 + 3 files changed, 929 insertions(+) create mode 100644 tensorrt_llm/_torch/models/modeling_beam.py 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 6d53abd9e065..c21259a417e4 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. @@ -626,6 +660,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..a53c8294eb8b --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_beam.py @@ -0,0 +1,848 @@ +# 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") + 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/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py index 2165c36576a5..b5182e70cfb1 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py +++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py @@ -128,6 +128,52 @@ def test_small_context_fallback_preserves_mixed_batch_generation( assert fallback.is_supported(q, None, None, metadata, forward_args, phase=phase) +@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) + 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( "dtype,sm_version", From 1e1caa12b990d37f926144b1dd685f1e7e53651e Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Sat, 17 Oct 2026 20:20:21 -0700 Subject: [PATCH 02/13] [None][feat] Add a Beam dense decoder-block component to ModelingV2 Add BeamDenseBlockBf16Sm103Tp1, a directly constructed component that runs Beam decoder block 0 (sliding-window attention, dense SwiGLU MLP) in BF16 on one sm_103 GPU, for one fresh context of 1 to 129 tokens. It is a component, not a routed target: there is no routing.py, router-index entry or ForCausalLM, and nothing in the engine builds it. The forward is assembled from catalog entries (CublasMm, FlashinferRmsnorm, FlashinferSiluAndMul, and ThopAttention over a KVCacheManagerV2 pool) and three entries the component owns: beam_neox_rope, beam_output_gate and beam_depth_residual. Configurations and inputs the component does not compute are rejected before any kernel runs. The catalog entries gain cells for the shapes the block runs. The block's tests compare the whole block, and every operation of it on its own input, against a plain-torch restatement of the source block at seeds {17, 41} x lengths {1, 17, 129}; they also check the op and kernel path and that a wrong residual coefficient is caught. They are local/QA tests and are not added to a CI test list. To run the tests: env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE \ python -m pytest tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1 Signed-off-by: Maxim Gluhovskoi --- .../activation/flashinfer_silu_and_mul.py | 10 +- .../catalog/attention/thop_attention.py | 32 +- .../modeling_v2/catalog/gemm/cublas_mm.py | 22 +- .../catalog/norm/flashinfer_rmsnorm.py | 20 +- .../modeling_v2/models/beam/__init__.py | 3 + .../dense_block_bf16__sm_103__tp1/__init__.py | 3 + .../dense_block_bf16__sm_103__tp1/block.py | 443 +++++++ .../ops/__init__.py | 3 + .../ops/attention/__init__.py | 2 + .../ops/attention/beam_neox_rope.py | 127 ++ .../ops/attention/beam_output_gate.py | 83 ++ .../ops/residual/__init__.py | 2 + .../ops/residual/beam_depth_residual.py | 92 ++ .../_dense_block_reference.py | 501 ++++++++ .../test_modeling_v2_beam_neox_rope.py | 83 ++ .../test_modeling_v2_beam_output_gate.py | 76 ++ .../test_modeling_v2_beam_depth_residual.py | 75 ++ .../test_dense_block.py | 1113 +++++++++++++++++ 18 files changed, 2686 insertions(+), 4 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/block.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_neox_rope.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_output_gate.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/beam_depth_residual.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/_dense_block_reference.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_neox_rope.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_output_gate.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/residual/test_modeling_v2_beam_depth_residual.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/test_dense_block.py 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 dca797613bff..7a6f0fe3ba89 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) + ), ) note = """ 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 7ec0b60bfe18..fa28d4e05301 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 @@ -63,6 +63,28 @@ mask_type=0, ) +#: 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, +) + #: 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( @@ -103,7 +125,8 @@ class ThopAttention(OpWrapper): # gpt-oss alternates its two window classes inside one engine, so both of # its cells are live every step; deepseek runs context and generation # against the same pool. qwen4_exp's cells give `cached` (tokens already in - # the pool) and `new` (this call's q rows) per sequence, contexts first. + # the pool) and `new` (this call's q rows) per sequence, contexts first, and + # so do the Beam dense block component's fresh-context cells. CELLS: tuple[Cell, ...] = ( Cell( why="gpt-oss sliding layer: GQA 64q/8kv d64, sinks, window 128", @@ -404,6 +427,13 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=3, ), ), + *( + 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) + ), ) note = """ 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 ce236e268841..3cc4db58f564 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,24 @@ 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) +) + class CublasMm(OpWrapper): """`mat_a @ mat_b (+ bias)`, fp32-accumulated, one cublas_mm call. @@ -24,7 +42,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 +63,7 @@ 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, ) note = """ 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 4a594aa80574..58c6911c001c 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/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/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..e66eac96a606 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/block.py @@ -0,0 +1,443 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam decoder block 0 / 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, loads the block's fifteen +tensors and runs one fresh BF16 context through it, with the attention batch state +prepared on a ``KVCacheManagerV2``. The forward is assembled from catalog entries and +this component's own ``ops/``, and it does not call the built-in Beam modeling. + +The block:: + + 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.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 + +from .ops.attention.beam_neox_rope import beam_neox_rope +from .ops.attention.beam_output_gate import beam_output_gate +from .ops.residual.beam_depth_residual import beam_depth_residual + +#: 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 +SUPPORTED_ROPE = {"rope_type": "default", "rope_theta": 10000.0} + +#: The longest fresh context the block was measured at; longer ones are rejected. +MAX_TOKENS = 129 + +_SM = (10, 3) + +#: thop_attention arguments that do not change with the step on this block: packed q|k|v, +#: append-and-attend, causal, bf16 pool, no quantization, and RoPE 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 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 layer_idx != SUPPORTED_LAYER_IDX: + raise ValueError( + f"only block {SUPPORTED_LAYER_IDX} is supported, got layer_idx={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 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: {layer_types[:1]!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: {mlp_types[:1]!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 decoder block 0: BF16 weights and activations, one GPU, one fresh context. + + ``forward(hidden_states [1, T, 5120] bf16, position_ids [1, T], attn_metadata)`` + returns ``[1, T, 5120]`` bf16 for 1 <= T <= 129, positions ``arange(T)`` and a + prepared ``TrtllmAttentionMetadata`` holding exactly that one context request on a + single-layer BF16 ``KVCacheManagerV2`` (8 kv heads of 128, no cached prefix). Inputs + outside that are rejected before any kernel runs; decode, cached prefixes, chunked + prefill, batches 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)) + + 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].""" + 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 tuple(tensor.shape) != shape or tensor.dtype != torch.bfloat16: + raise ValueError( + f"{name}: got {tuple(tensor.shape)} {tensor.dtype}, need {shape} torch.bfloat16" + ) + getattr(self, _attribute(name)).copy_(tensor) + + 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 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] != self.hidden_size + ): + raise ValueError( + f"hidden_states must be [1, T, {self.hidden_size}], got {tuple(hidden_states.shape)}" + ) + if hidden_states.device != self.input_layernorm.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 context lengths 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" + ) + expected = torch.arange(num_tokens, dtype=position_ids.dtype, device=position_ids.device) + if not torch.equal(position_ids[0], expected): + raise ValueError("a fresh context must have positions 0..T-1") + if not isinstance(md, TrtllmAttentionMetadata): + raise ValueError( + f"attn_metadata must be TrtllmAttentionMetadata, got {type(md).__name__}" + ) + mgr = md.kv_cache_manager + if mgr is None: + raise ValueError("attn_metadata needs a KV cache manager for the paged append") + if md.num_seqs != 1 or md.num_contexts != 1 or md.num_ctx_tokens != num_tokens: + raise ValueError( + f"one context request of {num_tokens} tokens is supported; metadata has " + f"{md.num_seqs} sequences, {md.num_contexts} contexts, {md.num_ctx_tokens} context tokens" + ) + if ( + int(md.kv_lens_runtime[0]) != num_tokens + or int(md.prompt_lens_cpu_runtime[0]) != num_tokens + ): + raise ValueError( + "a cached prefix is not supported: the request must start at an empty cache" + ) + if ( + md.use_paged_context_fmha + or md.is_cross + or md.is_spec_decoding_enabled + or md.beam_width != 1 + ): + raise ValueError( + "paged-context FMHA, cross attention, speculative decoding and beams are not supported" + ) + heads = getattr(mgr, "num_kv_heads_per_layer", None) or [None] + if ( + mgr.dtype != DataType.BF16 + or heads[self.layer_idx] != self.num_kv_heads + or mgr.head_dim != self.head_dim + ): + raise ValueError( + f"the KV cache must be BF16 with {self.num_kv_heads} kv heads of {self.head_dim}; " + f"got {mgr.dtype}, {heads}, head_dim {mgr.head_dim}" + ) + + @property + def hidden_size(self) -> int: + return self._shapes["input_layernorm.weight"][0] + + @staticmethod + def step_args(md: TrtllmAttentionMetadata) -> dict: + """Project the prepared metadata onto thop_attention's explicit batch state.""" + return dict( + 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(self, h: torch.Tensor, positions: torch.Tensor, step: dict) -> torch.Tensor: + """Normalized input [T, hidden] -> attention output [T, hidden], before post_attention_layernorm.""" + num_tokens = h.shape[0] + heads, kv_heads, d = self.num_heads, self.num_kv_heads, self.head_dim + q = cublas_mm(h, self.self_attn_q_proj.t()) + k = cublas_mm(h, self.self_attn_k_proj.t()) + v = cublas_mm(h, self.self_attn_v_proj.t()) + gate = cublas_mm(h, self.self_attn_gate_proj.t()) + q = flashinfer_rmsnorm(q.view(num_tokens * heads, d), self.self_attn_q_norm, self.eps) + k = flashinfer_rmsnorm(k.view(num_tokens * kv_heads, d), self.self_attn_k_norm, self.eps) + q, k = beam_neox_rope( + q.view(num_tokens, heads, d), + k.view(num_tokens, kv_heads, d), + positions, + self.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, + local_layer_idx=self.layer_idx, + num_heads=heads, + num_kv_heads=kv_heads, + head_size=d, + attention_window_size=self.window, + **step, + **_CALL_CONSTANTS, + ) + gated = beam_output_gate( + core.view(num_tokens, heads, d), + gate.view(num_tokens, heads, d), + self.self_attn_gate_bias, + ) + return cublas_mm(gated.view(num_tokens, heads * d), self.self_attn_o_proj.t()) + + 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/dense_block_bf16__sm_103__tp1/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_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/dense_block_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/dense_block_bf16__sm_103__tp1/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/__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/dense_block_bf16__sm_103__tp1/ops/attention/beam_neox_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_neox_rope.py new file mode 100644 index 000000000000..8faed2f17896 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_neox_rope.py @@ -0,0 +1,127 @@ +# 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. +_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/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_output_gate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_output_gate.py new file mode 100644 index 000000000000..1ee22c332098 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/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. +_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/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/__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/dense_block_bf16__sm_103__tp1/ops/residual/beam_depth_residual.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/beam_depth_residual.py new file mode 100644 index 000000000000..86317c5812d1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/beam_depth_residual.py @@ -0,0 +1,92 @@ +# 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 + +#: Block 0's two adds as (branch, residual) scales: (1, sqrt(n - 1)) / sqrt(n) for the +#: n-th add of the residual stream, n = 2 for the attention add and 3 for the feed-forward. +_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)) + +_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), + ), + ) + + 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/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/attention/test_modeling_v2_beam_neox_rope.py b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_neox_rope.py new file mode 100644 index 000000000000..d02865c44396 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/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.models.beam.dense_block_bf16__sm_103__tp1.ops.attention.beam_neox_rope import ( # noqa: E501 + 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/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_output_gate.py b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_output_gate.py new file mode 100644 index 000000000000..7d2b0030357e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/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.models.beam.dense_block_bf16__sm_103__tp1.ops.attention.beam_output_gate import ( # noqa: E501 + 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/models/beam/dense_block_bf16__sm_103__tp1/residual/test_modeling_v2_beam_depth_residual.py b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/residual/test_modeling_v2_beam_depth_residual.py new file mode 100644 index 000000000000..102b548557ae --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/residual/test_modeling_v2_beam_depth_residual.py @@ -0,0 +1,75 @@ +# 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.models.beam.dense_block_bf16__sm_103__tp1.ops.residual.beam_depth_residual import ( # noqa: E501 + 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 block 0's coefficients: the n-th residual add scales by + (1, sqrt(n - 1)) / sqrt(n), n = 2 for attention and 3 for the feed-forward.""" + 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)) + + +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) 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..c1240e9214e8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/test_dense_block.py @@ -0,0 +1,1113 @@ +# 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.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.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._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1.ops.attention.beam_neox_rope import ( # noqa: E501 + beam_neox_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1.ops.attention.beam_output_gate import ( # noqa: E501 + beam_output_gate, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1.ops.residual.beam_depth_residual import ( # noqa: E501 + beam_depth_residual, +) +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) + + 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, + local_layer_idx=LAYER_IDX, + num_heads=heads, + num_kv_heads=kv_heads, + head_size=d, + attention_window_size=BLOCK_CONFIG["sliding_window"], + **step, + **block_module._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 == 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: + with pytest.raises(ValueError, match="only block 0"): + BeamDenseBlockBf16Sm103Tp1( + block_config(layer_types=["sliding_attention"] * 2), 1, 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) + 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__), + "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) + 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) From 5b66179054aa0418a49deda443cd415670db1c0b Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Sun, 18 Oct 2026 16:15:21 -0700 Subject: [PATCH 03/13] [None][feat] Add a Beam MoE router component to ModelingV2 Add BeamMoeRouterBf16Sm103Tp1, a directly constructed component that runs the sigmoid grouped top-k router of Beam decoder block 3 on one sm_103 GPU. It takes already-normalized BF16 hidden states [T, 5120] and returns FP32 mixing weights [T, 8] and int64 expert ids [T, 8], in that order. It holds only the FP32 router weight [256, 5120] and the FP32 correction bias [256]. It does not run experts or build a MoE block, and nothing in the engine builds it. The arithmetic is one owned op, BeamGroupedTopkRouter, which follows the source router step by step in FP32 with TF32 and autocast off: projection, sigmoid, correction bias added for selection only, group score as the sum of the top two corrected scores in each of 8 groups of 32, 4 groups kept and the rest masked, top-8 experts, and the unbiased scores of the chosen experts divided by their sum plus 1e-20. The routed scaling factor is not applied. Tied scores keep the order torch.topk(sorted=True) returns; no index tie-break is added. The native noaux_tc op is not used because its sigmoid form, tie rule and scaling differ from the source. The tests run the component and compare every step of the op's call with the op's plain-torch reference at seeds {17, 41} x tokens {1, 17, 129}, and on controlled fixtures for the correction bias, the group and expert selection boundaries, group exclusion, ties and the epsilon. Group and expert ids are compared exactly and in order. They are local/QA tests and are not added to a CI test list. To run the tests: env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE \ python -m pytest tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1 Signed-off-by: Maxim Gluhovskoi --- .../moe_router_bf16__sm_103__tp1/__init__.py | 3 + .../ops/__init__.py | 3 + .../ops/moe/__init__.py | 2 + .../ops/moe/beam_grouped_topk_router.py | 191 ++++ .../moe_router_bf16__sm_103__tp1/router.py | 172 ++++ .../_router_cases.py | 420 ++++++++ ...st_modeling_v2_beam_grouped_topk_router.py | 59 ++ .../test_router.py | 912 ++++++++++++++++++ 8 files changed, 1762 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/ops/moe/beam_grouped_topk_router.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/router.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/_router_cases.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/moe/test_modeling_v2_beam_grouped_topk_router.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/test_router.py 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..6d45178b6936 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_router_bf16__sm_103__tp1/router.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE router 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 from the checkpoint's configuration, loads the router's two +tensors and passes it the already-normalized hidden states that enter block 3's +feed-forward sublayer. 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 + +#: 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 layer_idx != SUPPORTED_LAYER_IDX: + raise ValueError( + f"only block {SUPPORTED_LAYER_IDX}'s router is supported, got layer_idx={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 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): + """Block 3'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/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..1ff476a6ae63 --- /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 block 3"), + ], +) +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]) From 1b48e460aa3059c4dafdada16366bc4314eb3524 Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Mon, 19 Oct 2026 13:34:29 -0700 Subject: [PATCH 04/13] [None][feat] Add a Beam MoE routed-expert component to ModelingV2 Add BeamMoeExpertBf16Sm103Tp1, a directly constructed component that runs one routed SwiGLU expert of Beam decoder block 3 on one sm_103 GPU. It takes already-normalized BF16 hidden states [T, 5120], T >= 1, and returns the expert's unweighted BF16 output [T, 5120]. It holds only the three bias-free BF16 matrices gate_proj and up_proj [2560, 5120] and down_proj [5120, 2560], loaded by checkpoint name and copied as given. Every tensor is checked and staged on the device before any parameter changes, so a rejected load leaves the expert as it was. It routes no token and applies no mixing weight, routed scaling factor, residual or collective; nothing in the engine builds it. The three projections are the catalog's cublas_mm, each reading the stored weight's transpose as its column-major operand. The activation is one owned op, SwigluFp32: silu(gate.float()) * up.float(), rounded once to BF16 before the down projection, which also returns the FP32 SiLU and product so each step can be observed. The shared flashinfer_silu_and_mul entry is not used because it returns only the rounded activation. The tests run the component at seeds {17, 41} x tokens {1, 17, 129} and compare every step of the same forward call (gate, up, FP32 SiLU, FP32 product, BF16 activation, output) with the source expert written in plain torch, elementwise within one BF16 unit relative plus a per-step absolute allowance. Exact fixtures cover zero input, negative, zero and positive gates, a precision fixture on which every activation differs from one computed with a BF16 SiLU and product, asymmetric one-term feature maps that expose weight orientation and the gate/up/down mapping, token-row permutation, and zero down weights. They also check construction, exact loading, rejected inputs and loads, and that the forward runs on the GPU. The op's test drives its cells against an FP64 reference. They are local/QA tests and are not added to a CI test list. To run the tests: env -u TORCH_ALLOW_TF32_CUBLAS_OVERRIDE -u NVIDIA_TF32_OVERRIDE \ python -m pytest tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1 Signed-off-by: Maxim Gluhovskoi --- .../moe_expert_bf16__sm_103__tp1/__init__.py | 3 + .../moe_expert_bf16__sm_103__tp1/expert.py | 151 +++++++ .../ops/__init__.py | 3 + .../ops/activation/__init__.py | 2 + .../ops/activation/swiglu_fp32.py | 74 ++++ .../_expert_cases.py | 350 ++++++++++++++++ .../test_modeling_v2_swiglu_fp32.py | 81 ++++ .../test_expert.py | 380 ++++++++++++++++++ 8 files changed, 1044 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/expert.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/activation/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/ops/activation/swiglu_fp32.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/_expert_cases.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/activation/test_modeling_v2_swiglu_fp32.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_expert_bf16__sm_103__tp1/test_expert.py 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/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 From 1fba916fe82fc41d86f792408c65a39f971e4c9a Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Mon, 19 Oct 2026 18:20:09 -0700 Subject: [PATCH 05/13] [None][feat] Add a Beam MoE dispatch/combine component to ModelingV2 Add beam_dispatch_combine, a directly called component that dispatches the tokens of Beam decoder block 3 to their eight routed experts and combines the weighted expert outputs on one sm_103 GPU. It takes BF16 hidden states [T, 5120], T >= 1, the router's int64 expert ids and FP32 mixing weights [T, 8], and an expert provider, and returns BF16 [T, 5120]. It holds no state: no router parameter, no expert matrix and no configuration object. Nothing in the engine calls it. The provider owns the expert computation: provider(expert_id, rows) is called once per selected expert, in ascending id order, with that expert's BF16 rows in ascending token order, and returns the expert's unweighted BF16 output. An expert no token selected is never called. The component owns dispatch and combination. It checks its inputs and every provider output (type, dtype, shape, device) with errors naming the argument; the id range and per-row uniqueness are caller preconditions and are not scanned on the device. The arithmetic is one owned op, BeamDispatchCombine, written with ordinary torch CUDA operations: each weighted route is the FP32 product of the BF16 expert output and its FP32 weight, rounded once to BF16 and added to its token's FP32 accumulator in ascending expert order; the sum is multiplied by the routed scaling factor 2.5 once and rounded to BF16. The selected expert list and each expert's selection are host metadata; gathers, products, roundings, sums and the scaling run on the GPU. The tests run the public call at seeds {17, 41} x tokens {1, 17, 129} with controlled affine experts whose outputs depend on the expert id and the token values, and require bit-exact agreement with a plain-torch golden of the combination. Focused fixtures cover experts 0 and 255 shared by every token with unequal, unnormalized weights, a signed cancellation witness on which each rounding step and the single scale are visible, zero-weight routes, all-zero weights, all-zero expert outputs and token-row permutation. They also check rejected inputs and provider outputs, that the inputs are not modified, and that the call runs on the GPU. The op's test drives its cells bit-exactly against its own reference. The tests are collected by the existing l0_b300 pre-merge entry for unittest/_torch/modeling_v2: 51 tests, about 3 s on one GB300. To run the tests: python -m pytest tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1 Signed-off-by: Maxim Gluhovskoi --- .../__init__.py | 3 + .../dispatch_combine.py | 158 ++++++++ .../ops/__init__.py | 3 + .../ops/moe/__init__.py | 2 + .../ops/moe/beam_dispatch_combine.py | 132 +++++++ .../_dispatch_combine_cases.py | 226 ++++++++++++ .../test_modeling_v2_beam_dispatch_combine.py | 59 +++ .../test_dispatch_combine.py | 344 ++++++++++++++++++ 8 files changed, 927 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/dispatch_combine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/ops/moe/beam_dispatch_combine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/_dispatch_combine_cases.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/moe/test_modeling_v2_beam_dispatch_combine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_dispatch_combine_bf16__sm_103__tp1/test_dispatch_combine.py 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/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) From 7cf3acd9c414cffbbba305607985716c2d6c73db Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Tue, 20 Oct 2026 10:31:30 -0700 Subject: [PATCH 06/13] [None][feat] Add a Beam complete MoE and sparse decoder block to ModelingV2 Compose the accepted Beam MoE router, routed SwiGLU expert and dispatch/combine components into a directly constructed MoE feed-forward component, and build decoder block 3 (sliding attention, sparse MoE) on it. Both are components, not routed targets. - moe_bf16__sm_103__tp1: BeamMoeBf16Sm103Tp1 owns the FP32 router and 256 independently parameterized BF16 experts, checks all 770 tensors before loading any of them, and calls the accepted dispatch/combine with each expert as the provider. - sparse_block_bf16__sm_103__tp1: BeamSparseBlockBf16Sm103Tp1 runs the input RMSNorm, the sliding-window attention, the four sandwich norms, the MoE and block 3's depth-scaled residual adds. Its KV cache slot is the KVCacheManagerV2's local index for logical layer 3 (layer_offsets[3], selected with a layer_mask); a manager that does not hold layer 3 is rejected before any kernel runs. - _sliding_attention.py: the attention sublayer, step-argument projection and call checks shared by the dense and sparse blocks. The dense block delegates to it and resolves its cache slot the same way. - Promote beam_neox_rope, beam_output_gate and beam_depth_residual, with their tests, to the shared catalog, and add block 3's residual cells. Tests: generated-weight GPU tests for both components cover construction, loading, input and cache-metadata contracts, the six fixed cases against plain-torch references, designed routing fixtures, causal-prefix independence, residual coefficients and fault sensitivity. The dense block and moved operator tests are updated. To run the tests: python -m pytest tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1 \ tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1 Signed-off-by: Maxim Gluhovskoi --- .../attention/beam_neox_rope.py | 1 + .../attention/beam_output_gate.py | 2 +- .../residual}/__init__.py | 1 + .../residual/beam_depth_residual.py | 15 +- .../models/beam/_sliding_attention.py | 314 ++++ .../dense_block_bf16__sm_103__tp1/block.py | 253 +-- .../__init__.py | 1 + .../models/beam/moe_bf16__sm_103__tp1/moe.py | 210 +++ .../__init__.py | 2 +- .../sparse_block_bf16__sm_103__tp1/block.py | 293 ++++ .../test_modeling_v2_beam_neox_rope.py | 4 +- .../test_modeling_v2_beam_output_gate.py | 4 +- .../test_dense_block.py | 35 +- .../moe_bf16__sm_103__tp1/_moe_reference.py | 287 ++++ .../beam/moe_bf16__sm_103__tp1/test_moe.py | 1179 ++++++++++++++ .../_sparse_block_reference.py | 306 ++++ .../test_sparse_block.py | 1434 +++++++++++++++++ .../test_modeling_v2_beam_depth_residual.py | 15 +- 18 files changed, 4120 insertions(+), 236 deletions(-) rename tensorrt_llm/_torch/_experimental/modeling_v2/{models/beam/dense_block_bf16__sm_103__tp1/ops => catalog}/attention/beam_neox_rope.py (98%) rename tensorrt_llm/_torch/_experimental/modeling_v2/{models/beam/dense_block_bf16__sm_103__tp1/ops => catalog}/attention/beam_output_gate.py (97%) rename tensorrt_llm/_torch/_experimental/modeling_v2/{models/beam/dense_block_bf16__sm_103__tp1/ops/attention => catalog/residual}/__init__.py (85%) rename tensorrt_llm/_torch/_experimental/modeling_v2/{models/beam/dense_block_bf16__sm_103__tp1/ops => catalog}/residual/beam_depth_residual.py (81%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_sliding_attention.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/{dense_block_bf16__sm_103__tp1/ops/residual => moe_bf16__sm_103__tp1}/__init__.py (57%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/moe.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/{dense_block_bf16__sm_103__tp1/ops => sparse_block_bf16__sm_103__tp1}/__init__.py (64%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/block.py rename tests/unittest/_torch/modeling_v2/{models/beam/dense_block_bf16__sm_103__tp1 => }/attention/test_modeling_v2_beam_neox_rope.py (94%) rename tests/unittest/_torch/modeling_v2/{models/beam/dense_block_bf16__sm_103__tp1 => }/attention/test_modeling_v2_beam_output_gate.py (93%) create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1/_moe_reference.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/moe_bf16__sm_103__tp1/test_moe.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/_sparse_block_reference.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/test_sparse_block.py rename tests/unittest/_torch/modeling_v2/{models/beam/dense_block_bf16__sm_103__tp1 => }/residual/test_modeling_v2_beam_depth_residual.py (78%) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_neox_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_neox_rope.py similarity index 98% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_neox_rope.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_neox_rope.py index 8faed2f17896..4b5b952f3f71 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_neox_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_neox_rope.py @@ -9,6 +9,7 @@ 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 ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_output_gate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_output_gate.py similarity index 97% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_output_gate.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_output_gate.py index 1ee22c332098..5b87f8713ee1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/beam_output_gate.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_output_gate.py @@ -8,7 +8,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -#: Block 0's gate: 48 heads of 128, bf16. +#: 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) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/__init__.py similarity index 85% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/__init__.py index 52a7a9daf028..51c47b7ad054 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/attention/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/__init__.py @@ -1,2 +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/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/beam_depth_residual.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/beam_depth_residual.py similarity index 81% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/beam_depth_residual.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/beam_depth_residual.py index 86317c5812d1..4706d71f2de8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/beam_depth_residual.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/beam_depth_residual.py @@ -10,10 +10,13 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -#: Block 0's two adds as (branch, residual) scales: (1, sqrt(n - 1)) / sqrt(n) for the -#: n-th add of the residual stream, n = 2 for the attention add and 3 for the feed-forward. +#: 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) @@ -41,6 +44,14 @@ class BeamDepthResidual(OpWrapper): 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 = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_sliding_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_sliding_attention.py new file mode 100644 index 000000000000..e0fff7ad450c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_sliding_attention.py @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sliding-window attention sublayer of the Beam decoder block components, one fresh context. + +Shared by the dense block 0 and the sparse block 3 components. It holds no parameters, +cache or configuration: each block passes its own tensors and geometry 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 + 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 + out = a @ o_proj^T + +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 feed-forward kind). The KV cache it +appends to is a ``KVCacheManagerV2`` that holds only 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 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 position check compares the positions with ``arange(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_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 RoPE 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 local_layer_index(md: TrtllmAttentionMetadata, layer_idx: int) -> int: + """The cache slot of logical layer ``layer_idx`` in the call's single-layer cache manager. + + Raises ValueError unless the manager holds exactly one local layer and that layer is + ``layer_idx``. 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 num_local != 1 or not 0 <= local < num_local: + raise ValueError( + f"a single-layer KV cache is supported: the manager has {num_local} local layers " + f"and maps logical layer {layer_idx} to slot {local}" + ) + 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(T)`` + on that device; a prepared ``TrtllmAttentionMetadata`` holding one fresh context + request of T tokens on a BF16 ``KVCacheManagerV2`` that holds exactly logical layer + ``layer_idx`` with ``num_kv_heads`` kv heads of ``head_dim``. + """ + 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 context lengths 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" + ) + expected = torch.arange(num_tokens, dtype=position_ids.dtype, device=position_ids.device) + if not torch.equal(position_ids[0], expected): + raise ValueError("a fresh context must have positions 0..T-1") + if not isinstance(md, TrtllmAttentionMetadata): + raise ValueError(f"attn_metadata must be TrtllmAttentionMetadata, got {type(md).__name__}") + mgr = md.kv_cache_manager + if mgr is None: + raise ValueError("attn_metadata needs a KV cache manager for the paged append") + if md.num_seqs != 1 or md.num_contexts != 1 or md.num_ctx_tokens != num_tokens: + raise ValueError( + f"one context request of {num_tokens} tokens is supported; metadata has " + f"{md.num_seqs} sequences, {md.num_contexts} contexts, {md.num_ctx_tokens} context tokens" + ) + if int(md.kv_lens_runtime[0]) != num_tokens or int(md.prompt_lens_cpu_runtime[0]) != num_tokens: + raise ValueError( + "a cached prefix is not supported: the request must start at an empty cache" + ) + if ( + md.use_paged_context_fmha + or md.is_cross + or md.is_spec_decoding_enabled + or md.beam_width != 1 + ): + raise ValueError( + "paged-context FMHA, cross attention, speculative decoding and beams are not supported" + ) + 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}" + ) + + +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, + 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]``. + """ + 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 = beam_neox_rope( + q.view(num_tokens, heads, d), k.view(num_tokens, kv_heads, d), 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/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 index e66eac96a606..e76074aa7a41 100644 --- 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 @@ -6,8 +6,9 @@ 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 block's fifteen tensors and runs one fresh BF16 context through it, with the attention batch state -prepared on a ``KVCacheManagerV2``. The forward is assembled from catalog entries and -this component's own ``ops/``, and it does not call the built-in Beam modeling. +prepared on a ``KVCacheManagerV2``. The forward is assembled from catalog entries, with +the attention sublayer shared with the sparse block component in ``../_sliding_attention.py``, +and it does not call the built-in Beam modeling. The block:: @@ -37,20 +38,16 @@ 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.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._experimental.modeling_v2.catalog.residual.beam_depth_residual import ( + beam_depth_residual, +) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata -from tensorrt_llm.bindings import DataType -from .ops.attention.beam_neox_rope import beam_neox_rope -from .ops.attention.beam_output_gate import beam_output_gate -from .ops.residual.beam_depth_residual import beam_depth_residual +from .. import _sliding_attention #: The configuration this component computes, as the checkpoint's config.json states it. #: Anything else is rejected at construction rather than computed approximately. @@ -74,64 +71,6 @@ _SM = (10, 3) -#: thop_attention arguments that do not change with the step on this block: packed q|k|v, -#: append-and-attend, causal, bf16 pool, no quantization, and RoPE 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 weight_shapes(config) -> "OrderedDict[str, tuple[int, ...]]": """The block's fifteen tensors by checkpoint name, in checkpoint order, [out, in] for matrices.""" @@ -206,9 +145,10 @@ class BeamDenseBlockBf16Sm103Tp1(nn.Module): ``forward(hidden_states [1, T, 5120] bf16, position_ids [1, T], attn_metadata)`` returns ``[1, T, 5120]`` bf16 for 1 <= T <= 129, positions ``arange(T)`` and a prepared ``TrtllmAttentionMetadata`` holding exactly that one context request on a - single-layer BF16 ``KVCacheManagerV2`` (8 kv heads of 128, no cached prefix). Inputs - outside that are rejected before any kernel runs; decode, cached prefixes, chunked - prefill, batches and other blocks are not computed here. + single-layer BF16 ``KVCacheManagerV2`` (8 kv heads of 128, no cached prefix) that + holds logical layer 0. Inputs outside that are rejected before any kernel runs; + decode, cached prefixes, chunked prefill, batches and other blocks are not computed + here. """ def __init__( @@ -258,158 +198,47 @@ 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 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] != self.hidden_size - ): - raise ValueError( - f"hidden_states must be [1, T, {self.hidden_size}], got {tuple(hidden_states.shape)}" - ) - if hidden_states.device != self.input_layernorm.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 context lengths 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" - ) - expected = torch.arange(num_tokens, dtype=position_ids.dtype, device=position_ids.device) - if not torch.equal(position_ids[0], expected): - raise ValueError("a fresh context must have positions 0..T-1") - if not isinstance(md, TrtllmAttentionMetadata): - raise ValueError( - f"attn_metadata must be TrtllmAttentionMetadata, got {type(md).__name__}" - ) - mgr = md.kv_cache_manager - if mgr is None: - raise ValueError("attn_metadata needs a KV cache manager for the paged append") - if md.num_seqs != 1 or md.num_contexts != 1 or md.num_ctx_tokens != num_tokens: - raise ValueError( - f"one context request of {num_tokens} tokens is supported; metadata has " - f"{md.num_seqs} sequences, {md.num_contexts} contexts, {md.num_ctx_tokens} context tokens" - ) - if ( - int(md.kv_lens_runtime[0]) != num_tokens - or int(md.prompt_lens_cpu_runtime[0]) != num_tokens - ): - raise ValueError( - "a cached prefix is not supported: the request must start at an empty cache" - ) - if ( - md.use_paged_context_fmha - or md.is_cross - or md.is_spec_decoding_enabled - or md.beam_width != 1 - ): - raise ValueError( - "paged-context FMHA, cross attention, speculative decoding and beams are not supported" - ) - heads = getattr(mgr, "num_kv_heads_per_layer", None) or [None] - if ( - mgr.dtype != DataType.BF16 - or heads[self.layer_idx] != self.num_kv_heads - or mgr.head_dim != self.head_dim - ): - raise ValueError( - f"the KV cache must be BF16 with {self.num_kv_heads} kv heads of {self.head_dim}; " - f"got {mgr.dtype}, {heads}, head_dim {mgr.head_dim}" - ) + _sliding_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, + ) @property def hidden_size(self) -> int: return self._shapes["input_layernorm.weight"][0] - @staticmethod - def step_args(md: TrtllmAttentionMetadata) -> dict: + def step_args(self, md: TrtllmAttentionMetadata) -> dict: """Project the prepared metadata onto thop_attention's explicit batch state.""" - return dict( - 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, - ) + return _sliding_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.""" - num_tokens = h.shape[0] - heads, kv_heads, d = self.num_heads, self.num_kv_heads, self.head_dim - q = cublas_mm(h, self.self_attn_q_proj.t()) - k = cublas_mm(h, self.self_attn_k_proj.t()) - v = cublas_mm(h, self.self_attn_v_proj.t()) - gate = cublas_mm(h, self.self_attn_gate_proj.t()) - q = flashinfer_rmsnorm(q.view(num_tokens * heads, d), self.self_attn_q_norm, self.eps) - k = flashinfer_rmsnorm(k.view(num_tokens * kv_heads, d), self.self_attn_k_norm, self.eps) - q, k = beam_neox_rope( - q.view(num_tokens, heads, d), - k.view(num_tokens, kv_heads, d), + return _sliding_attention.attention( + h, positions, - self.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, - local_layer_idx=self.layer_idx, - num_heads=heads, - num_kv_heads=kv_heads, - head_size=d, - attention_window_size=self.window, - **step, - **_CALL_CONSTANTS, - ) - gated = beam_output_gate( - core.view(num_tokens, heads, d), - gate.view(num_tokens, heads, d), - self.self_attn_gate_bias, + 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, ) - return cublas_mm(gated.view(num_tokens, heads * d), self.self_attn_o_proj.t()) def feed_forward(self, h: torch.Tensor) -> torch.Tensor: """Normalized input [T, hidden] -> dense SwiGLU output [T, hidden], before its post-norm.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/__init__.py similarity index 57% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/__init__.py index 52a7a9daf028..e6278158bae4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/residual/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/__init__.py @@ -1,2 +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..7164afaf120a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/moe_bf16__sm_103__tp1/moe.py @@ -0,0 +1,210 @@ +# 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, 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 block 3's feed-forward +sublayer. 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): + """Block 3'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/dense_block_bf16__sm_103__tp1/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/__init__.py similarity index 64% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/__init__.py index 23e614d4422b..c10dd6cb1a2e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/ops/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/__init__.py @@ -1,3 +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.""" +"""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..9934aec4ae9d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_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 block 3 / sm_103 / tp1: one sparse, 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, loads the block's 782 tensors +and runs one fresh BF16 context through it, with the attention batch state prepared on a +``KVCacheManagerV2`` that holds logical layer 3. The forward is assembled from catalog +entries, the attention sublayer shared with the dense block component +(``../_sliding_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:: + + h = rmsnorm(x, input_layernorm) + a = sliding attention(h) see ../_sliding_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) + +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 logical block 3'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`` = 3 fixes the block's semantics; the cache +slot is the call's manager's local index for layer 3, 0 for a manager whose +``layer_mask`` selects only layer 3. A manager that does not hold layer 3 -- 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). Lengths up to 129 tokens never +reach the 512-token window, so behaviour at or beyond the window is not exercised here. +""" + +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 _sliding_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 +SUPPORTED_ROPE = {"rope_type": "default", "rope_theta": 10000.0} + +#: The longest fresh context the block was measured at; longer ones are rejected. +MAX_TOKENS = 129 + +_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 layer_idx != SUPPORTED_LAYER_IDX: + raise ValueError( + f"only block {SUPPORTED_LAYER_IDX} is supported, got layer_idx={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 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}" + ) + 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}" + ) + check_moe_config(config, layer_idx) + + +class BeamSparseBlockBf16Sm103Tp1(nn.Module): + """Beam decoder block 3: BF16 weights and activations, an FP32 router, one GPU, one fresh context. + + 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 <= 129, positions ``arange(T)`` and a + prepared ``TrtllmAttentionMetadata`` holding exactly that one context request on a + BF16 ``KVCacheManagerV2`` (8 kv heads of 128, no cached prefix) that holds logical + layer 3 and nothing else. Inputs outside that are rejected before any kernel runs; + decode, cached prefixes, chunked prefill, batches 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.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.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") + _sliding_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 _sliding_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 _sliding_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, + ) + + @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/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_neox_rope.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_neox_rope.py similarity index 94% rename from tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_neox_rope.py rename to tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_neox_rope.py index d02865c44396..3eef5075bb85 100644 --- a/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_neox_rope.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_neox_rope.py @@ -7,11 +7,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1.ops.attention.beam_neox_rope import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_neox_rope import ( beam_neox_rope as op, ) -__extra_import_path__ = ["../../../.."] +__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" diff --git a/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_output_gate.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_output_gate.py similarity index 93% rename from tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_output_gate.py rename to tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_output_gate.py index 7d2b0030357e..4a68cc0433bf 100644 --- a/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/attention/test_modeling_v2_beam_output_gate.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_output_gate.py @@ -7,11 +7,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1.ops.attention.beam_output_gate import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.beam_output_gate import ( beam_output_gate as op, ) -__extra_import_path__ = ["../../../.."] +__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" 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 index c1240e9214e8..2ed904ecba02 100644 --- 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 @@ -56,6 +56,12 @@ 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, ) @@ -63,21 +69,16 @@ 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 _sliding_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._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1.ops.attention.beam_neox_rope import ( # noqa: E501 - beam_neox_rope, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1.ops.attention.beam_output_gate import ( # noqa: E501 - beam_output_gate, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1.ops.residual.beam_depth_residual import ( # noqa: E501 - beam_depth_residual, -) 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 @@ -493,6 +494,8 @@ def replay( 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. @@ -551,13 +554,12 @@ def check(name: str, rule: str | None, actual: torch.Tensor, expected: str) -> d core = thop_attention( q=qkv, output=core, - local_layer_idx=LAYER_IDX, num_heads=heads, num_kv_heads=kv_heads, head_size=d, attention_window_size=BLOCK_CONFIG["sliding_window"], **step, - **block_module._CALL_CONSTANTS, + **_sliding_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")) @@ -1031,6 +1033,12 @@ def test_forward_runs_the_component_path(seed, length, ieee_fp32) -> None: 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}) @@ -1047,7 +1055,7 @@ def test_forward_runs_the_component_path(seed, length, ieee_fp32) -> None: "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__), + "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) @@ -1068,6 +1076,9 @@ def test_forward_runs_the_component_path(seed, length, ieee_fp32) -> None: 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 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..5b074b1bb50a --- /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 block 3"), +] + + +@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/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..6e149331fe18 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/test_sparse_block.py @@ -0,0 +1,1434 @@ +# 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 == 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, 5): + with pytest.raises(ValueError, match="only block 3"): + 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", + "two_layer_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) and one that holds layer 3 beside another layer.""" + 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) + elif case == "two_layer_cache": + context = dict( + num_tokens=length, num_layers=2, layer_mask=[i in (3, 4) for i in range(NUM_LAYERS)] + ) + 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} + elif case == "two_layer_cache": + assert dict(md.kv_cache_manager.layer_offsets) == {3: 0, 4: 1} + 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) + + +# -------------------------------------------------- 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/dense_block_bf16__sm_103__tp1/residual/test_modeling_v2_beam_depth_residual.py b/tests/unittest/_torch/modeling_v2/residual/test_modeling_v2_beam_depth_residual.py similarity index 78% rename from tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/residual/test_modeling_v2_beam_depth_residual.py rename to tests/unittest/_torch/modeling_v2/residual/test_modeling_v2_beam_depth_residual.py index 102b548557ae..670b42253440 100644 --- a/tests/unittest/_torch/modeling_v2/models/beam/dense_block_bf16__sm_103__tp1/residual/test_modeling_v2_beam_depth_residual.py +++ b/tests/unittest/_torch/modeling_v2/residual/test_modeling_v2_beam_depth_residual.py @@ -8,11 +8,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.beam.dense_block_bf16__sm_103__tp1.ops.residual.beam_depth_residual import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.residual.beam_depth_residual import ( beam_depth_residual as op, ) -__extra_import_path__ = ["../../../.."] +__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" @@ -39,11 +39,18 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_scales_are_the_source_depth_coefficients() -> None: - """The cells carry block 0's coefficients: the n-th residual add scales by - (1, sqrt(n - 1)) / sqrt(n), n = 2 for attention and 3 for the feed-forward.""" + """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: From 245fb0b1e69e197e573f751a3531d10ee107c13e Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Tue, 20 Oct 2026 19:20:30 -0700 Subject: [PATCH 07/13] [None][feat] Add a Beam six-layer prefill model to ModelingV2 Compose embeddings, three dense and three sparse layers, final normalization, an untied output head and softcap using generated weights. Generalize shared attention for logical layer slots and full-context NoPE query scaling. Validated eager TP1 on GB300 for seeds 17/41 and lengths 1/17/129: 80 public checks and 354 component regressions passed, with independent review and QA acceptance. Scope is six-layer fresh-context correctness; cache reuse, decoding, real weights and full-depth inference remain future work. Signed-off-by: Maxim Gluhovskoi --- .../attention/beam_nope_query_scale.py | 64 ++ .../modeling_v2/catalog/gemm/cublas_mm.py | 9 + .../{_sliding_attention.py => _attention.py} | 91 +- .../dense_block_bf16__sm_103__tp1/block.py | 46 +- .../models/beam/moe_bf16__sm_103__tp1/moe.py | 10 +- .../moe_router_bf16__sm_103__tp1/router.py | 21 +- .../__init__.py | 3 + .../model.py | 508 ++++++++++++ .../ops/__init__.py | 3 + .../ops/activation/__init__.py | 2 + .../ops/activation/beam_logit_softcap.py | 61 ++ .../sparse_block_bf16__sm_103__tp1/block.py | 77 +- .../test_modeling_v2_beam_nope_query_scale.py | 85 ++ .../test_dense_block.py | 11 +- .../beam/moe_bf16__sm_103__tp1/test_moe.py | 2 +- .../test_router.py | 2 +- .../test_beam_logit_softcap.py | 76 ++ .../test_small_model_prefill.py | 777 ++++++++++++++++++ .../test_sparse_block.py | 39 +- 19 files changed, 1779 insertions(+), 108 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/beam_nope_query_scale.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/{_sliding_attention.py => _attention.py} (76%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/model.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/activation/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/ops/activation/beam_logit_softcap.py create mode 100644 tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_beam_nope_query_scale.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_beam_logit_softcap.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_small_model_prefill.py 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/gemm/cublas_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py index 3cc4db58f564..dff9afec2ebd 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 @@ -27,6 +27,14 @@ ) 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): @@ -64,6 +72,7 @@ class CublasMm(OpWrapper): spec=dict(m=1, k=7168, n=1536, dtype=torch.bfloat16, bias=False, out_dtype=None), ), *_BEAM_DENSE_BLOCK, + *_BEAM_HEAD, ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_sliding_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py similarity index 76% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_sliding_attention.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py index e0fff7ad450c..8e25b9519916 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_sliding_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py @@ -1,26 +1,33 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The sliding-window attention sublayer of the Beam decoder block components, one fresh context. +"""The attention sublayer of the Beam decoder block components, one fresh context. -Shared by the dense block 0 and the sparse block 3 components. It holds no parameters, -cache or configuration: each block passes its own tensors and geometry and the prepared -attention metadata of the call. The computation, before the block's post-attention norm:: +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 - 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 + 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. + 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 feed-forward kind). The KV cache it -appends to is a ``KVCacheManagerV2`` that holds only 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 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. +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 position check compares the positions with ``arange(T)`` on the device and reads the answer back, once per call. The metadata checks read host tensors only. @@ -35,6 +42,9 @@ 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, ) @@ -50,8 +60,9 @@ 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 RoPE applied before the -#: call -- the rope_* values are inert placeholders the op requires, not this model's RoPE. +#: 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, @@ -109,10 +120,10 @@ def local_layer_index(md: TrtllmAttentionMetadata, layer_idx: int) -> int: - """The cache slot of logical layer ``layer_idx`` in the call's single-layer cache manager. + """The cache slot of logical layer ``layer_idx`` in the call's cache manager. - Raises ValueError unless the manager holds exactly one local layer and that layer is - ``layer_idx``. Reads the manager's host-side layer map only. + 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) @@ -124,10 +135,15 @@ def local_layer_index(md: TrtllmAttentionMetadata, layer_idx: int) -> int: ) num_local = getattr(mgr, "num_local_layers", None) local = offsets[layer_idx] - if num_local != 1 or not 0 <= local < num_local: + 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"a single-layer KV cache is supported: the manager has {num_local} local layers " - f"and maps logical layer {layer_idx} to slot {local}" + 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 @@ -149,7 +165,7 @@ def check_step( Accepted: ``hidden_states`` contiguous ``[1, T, hidden_size]`` bf16 on ``device``, 1 <= T <= ``max_tokens``; integer ``position_ids`` ``[1, T]`` equal to ``arange(T)`` on that device; a prepared ``TrtllmAttentionMetadata`` holding one fresh context - request of T tokens on a BF16 ``KVCacheManagerV2`` that holds exactly logical layer + request of T tokens on a BF16 ``KVCacheManagerV2`` that holds logical layer ``layer_idx`` with ``num_kv_heads`` kv heads of ``head_dim``. """ if not isinstance(hidden_states, torch.Tensor) or hidden_states.dtype != torch.bfloat16: @@ -179,6 +195,23 @@ def check_step( expected = torch.arange(num_tokens, dtype=position_ids.dtype, device=position_ids.device) if not torch.equal(position_ids[0], expected): raise ValueError("a fresh context must have positions 0..T-1") + check_metadata( + md, num_tokens, layer_idx=layer_idx, num_kv_heads=num_kv_heads, head_dim=head_dim + ) + + +def check_metadata( + md: TrtllmAttentionMetadata, + num_tokens: int, + *, + layer_idx: int, + num_kv_heads: int, + head_dim: int, +) -> None: + """Raise ValueError unless ``md`` is one prepared fresh context request of ``num_tokens`` + tokens on a BF16 ``KVCacheManagerV2`` holding logical layer ``layer_idx`` with + ``num_kv_heads`` kv heads of ``head_dim``. Reads host-side state only. + """ if not isinstance(md, TrtllmAttentionMetadata): raise ValueError(f"attn_metadata must be TrtllmAttentionMetadata, got {type(md).__name__}") mgr = md.kv_cache_manager @@ -278,12 +311,16 @@ def attention( num_kv_heads: int, head_dim: int, eps: float, - rope_theta: 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 @@ -293,9 +330,11 @@ def attention( 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 = beam_neox_rope( - q.view(num_tokens, heads, d), k.view(num_tokens, kv_heads, d), positions, rope_theta, 1.0 - ) + 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( 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 index e76074aa7a41..92b1ba87658e 100644 --- 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 @@ -1,16 +1,16 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Beam decoder block 0 / sm_103 / tp1: one dense, sliding-attention block, constructed directly. +"""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, loads the block's fifteen -tensors and runs one fresh BF16 context through it, 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 ``../_sliding_attention.py``, -and it does not call the built-in Beam modeling. +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 fresh BF16 context through it, +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:: +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 @@ -47,7 +47,7 @@ ) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata -from .. import _sliding_attention +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. @@ -64,6 +64,8 @@ "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 longest fresh context the block was measured at; longer ones are rejected. @@ -116,9 +118,9 @@ def depth_scales(layer_idx: int, is_attention: bool) -> tuple[float, float]: def check_config(config, layer_idx: int) -> None: """Raise ValueError naming the first field outside what this component computes.""" - if layer_idx != SUPPORTED_LAYER_IDX: + if type(layer_idx) is not int or layer_idx not in SUPPORTED_LAYER_INDICES: raise ValueError( - f"only block {SUPPORTED_LAYER_IDX} is supported, got layer_idx={layer_idx}" + 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) @@ -127,11 +129,15 @@ def check_config(config, layer_idx: int) -> None: 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: {layer_types[:1]!r}" + 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: {mlp_types[:1]!r}") + 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( @@ -140,15 +146,15 @@ def check_config(config, layer_idx: int) -> None: class BeamDenseBlockBf16Sm103Tp1(nn.Module): - """Beam decoder block 0: BF16 weights and activations, one GPU, one fresh context. + """Beam dense decoder block 0, 1 or 2: BF16 weights and activations, one GPU, one fresh context. ``forward(hidden_states [1, T, 5120] bf16, position_ids [1, T], attn_metadata)`` returns ``[1, T, 5120]`` bf16 for 1 <= T <= 129, positions ``arange(T)`` and a prepared ``TrtllmAttentionMetadata`` holding exactly that one context request on a - single-layer BF16 ``KVCacheManagerV2`` (8 kv heads of 128, no cached prefix) that - holds logical layer 0. Inputs outside that are rejected before any kernel runs; - decode, cached prefixes, chunked prefill, batches and other blocks are not computed - here. + BF16 ``KVCacheManagerV2`` (8 kv heads of 128, no cached prefix) that holds the + block's logical layer ``layer_idx``. Inputs outside that are rejected before any + kernel runs; decode, cached prefixes, chunked prefill, batches and other blocks are + not computed here. """ def __init__( @@ -198,7 +204,7 @@ 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.""" - _sliding_attention.check_step( + _attention.check_step( hidden_states, position_ids, md, @@ -216,11 +222,11 @@ def hidden_size(self) -> int: def step_args(self, md: TrtllmAttentionMetadata) -> dict: """Project the prepared metadata onto thop_attention's explicit batch state.""" - return _sliding_attention.step_args(md, self.layer_idx) + 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 _sliding_attention.attention( + return _attention.attention( h, positions, step, 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 index 7164afaf120a..626143c4acda 100644 --- 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 @@ -1,12 +1,14 @@ # 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, constructed directly. +"""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 block 3's feed-forward -sublayer. It composes three accepted components and adds no arithmetic of its own:: +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 @@ -103,7 +105,7 @@ def check_weights(expected: Mapping, weights: Mapping[str, torch.Tensor], what: class BeamMoeBf16Sm103Tp1(nn.Module): - """Block 3's MoE: normalized BF16 hidden states in, the combined BF16 MoE output out. + """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 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 index 6d45178b6936..8a88d7d0309f 100644 --- 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 @@ -1,15 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Beam MoE router of decoder block 3 / sm_103 / tp1, constructed directly. +"""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 block 3's -feed-forward sublayer. 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. +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``):: @@ -39,6 +39,8 @@ "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 @@ -48,9 +50,10 @@ def check_config(config, layer_idx: int) -> None: """Raise ValueError naming the first field outside what this component computes.""" - if layer_idx != SUPPORTED_LAYER_IDX: + if type(layer_idx) is not int or layer_idx not in SUPPORTED_LAYER_INDICES: raise ValueError( - f"only block {SUPPORTED_LAYER_IDX}'s router is supported, got layer_idx={layer_idx}" + 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) @@ -66,7 +69,7 @@ def check_config(config, layer_idx: int) -> None: class BeamMoeRouterBf16Sm103Tp1(nn.Module): - """Block 3's router: BF16 hidden states in, expert ids and FP32 mixing weights out. + """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`` 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/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..e6f527d6e382 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/model.py @@ -0,0 +1,508 @@ +# 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 prefill logits, constructed directly. + +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. 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 fresh BF16 context through it, +with the attention batch state prepared on a six-layer ``KVCacheManagerV2`` +(``prefill_context`` builds one). 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. + +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 keep a 512-token window and layer 5 keeps the whole context; 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. ``check_cache`` verifies all of it before the first layer runs, including that +the host pool tables the attention op reads address each slot's own key and value pages +as the manager allocated them. Cache contents are written as the backend requires; they +are not read back, reused or shared between calls. + +Host synchronization per call: one for the model's token-id range and position 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 +from collections import OrderedDict +from typing import Iterator, Mapping + +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 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 longest fresh context the model was measured at; longer ones are rejected. +MAX_TOKENS = 129 + +#: 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"],) + +_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, num_tokens: 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 call's tokens; keys and values in one shared page + table per pool; and 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``). Reads host-side manager state and host tensors 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 < num_tokens 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 call has {num_tokens}" + ) + _check_pool_tables(md, mgr, slots) + + +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 token t's + softcapped logit for vocabulary id v. Accepted: 1 <= T <= 129 token ids in + ``[0, 151936)``, contiguous on the model's device; integer positions equal to + ``arange(T)``; and the prepared ``TrtllmAttentionMetadata`` of one fresh context + request of T tokens on the six-layer cache ``check_cache`` describes, such as + ``prefill_context(T)`` yields. 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. Decode, cached prefixes, chunked prefill, + batches, padding and contexts longer than 129 tokens are not computed here. + """ + + 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 + + @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. + """ + 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 _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 context lengths 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__}" + ) + check_cache(md, num_tokens) + for layer_idx in range(NUM_LAYERS): + _attention.check_metadata( + md, num_tokens, layer_idx=layer_idx, num_kv_heads=8, head_dim=128 + ) + expected = torch.arange(num_tokens, dtype=position_ids.dtype, device=self.device) + low, high = torch.aminmax(input_ids) + positions_ok, ids_ok = torch.stack( + ((position_ids[0] == expected).all(), (low >= 0) & (high < self.vocab_size)) + ).tolist() + if not positions_ok: + raise ValueError("a fresh context must have positions 0..T-1") + 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/block.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/block.py index 9934aec4ae9d..325f763c51d1 100644 --- 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 @@ -1,34 +1,40 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Beam decoder block 3 / sm_103 / tp1: one sparse, sliding-attention block, constructed directly. +"""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, loads the block's 782 tensors -and runs one fresh BF16 context through it, with the attention batch state prepared on a -``KVCacheManagerV2`` that holds logical layer 3. The forward is assembled from catalog -entries, the attention sublayer shared with the dense block component -(``../_sliding_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. +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 fresh BF16 context through it, 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:: +The block, with block 3's residual coefficients (block i's are in ``depth_scales``):: h = rmsnorm(x, input_layernorm) - a = sliding attention(h) see ../_sliding_attention.py + 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 logical block 3's: block i's attention add is the 2(i + 1)-th add of +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`` = 3 fixes the block's semantics; the cache -slot is the call's manager's local index for layer 3, 0 for a manager whose -``layer_mask`` selects only layer 3. A manager that does not hold layer 3 -- a default -one-layer manager holds layer 0 -- is rejected before any kernel runs. +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). Lengths up to 129 tokens never @@ -52,7 +58,7 @@ ) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata -from .. import _sliding_attention +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 @@ -70,6 +76,10 @@ "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 longest fresh context the block was measured at; longer ones are rejected. @@ -122,30 +132,34 @@ def depth_scales(layer_idx: int, is_attention: bool) -> tuple[float, float]: def check_config(config, layer_idx: int) -> None: """Raise ValueError naming the first field outside what this component computes.""" - if layer_idx != SUPPORTED_LAYER_IDX: + if type(layer_idx) is not int or layer_idx not in SUPPORTED_LAYER_INDICES: raise ValueError( - f"only block {SUPPORTED_LAYER_IDX} is supported, got layer_idx={layer_idx}" + 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] != "sliding_attention": + if len(layer_types) <= layer_idx or layer_types[layer_idx] != attention_type: raise ValueError( - f"block {layer_idx} must be a sliding_attention layer: " + 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 = (getattr(config, "rope_parameters", None) or {}).get("sliding_attention") - if rope != SUPPORTED_ROPE: + 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 sliding_attention rope parameters {rope!r}; need {SUPPORTED_ROPE!r}" + f"unsupported {attention_type} rope parameters {rope!r}; need {expected_rope!r}" ) check_moe_config(config, layer_idx) class BeamSparseBlockBf16Sm103Tp1(nn.Module): - """Beam decoder block 3: BF16 weights and activations, an FP32 router, one GPU, one fresh context. + """Beam sparse decoder block 3, 4 or 5: BF16 weights and activations, an FP32 router, one GPU, + one fresh context. 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. @@ -153,8 +167,8 @@ class BeamSparseBlockBf16Sm103Tp1(nn.Module): ``forward(hidden_states [1, T, 5120] bf16, position_ids [1, T], attn_metadata)`` returns ``[1, T, 5120]`` bf16 for 1 <= T <= 129, positions ``arange(T)`` and a prepared ``TrtllmAttentionMetadata`` holding exactly that one context request on a - BF16 ``KVCacheManagerV2`` (8 kv heads of 128, no cached prefix) that holds logical - layer 3 and nothing else. Inputs outside that are rejected before any kernel runs; + BF16 ``KVCacheManagerV2`` (8 kv heads of 128, no cached prefix) that holds the + block's logical layer. Inputs outside that are rejected before any kernel runs; decode, cached prefixes, chunked prefill, batches and other blocks are not computed here. Inputs and parameters are never modified. """ @@ -174,8 +188,9 @@ def __init__( 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.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) @@ -230,7 +245,7 @@ def _check_step( """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") - _sliding_attention.check_step( + _attention.check_step( hidden_states, position_ids, md, @@ -244,11 +259,11 @@ def _check_step( def step_args(self, md: TrtllmAttentionMetadata) -> dict: """Project the prepared metadata onto thop_attention's explicit batch state.""" - return _sliding_attention.step_args(md, self.layer_idx) + 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 _sliding_attention.attention( + return _attention.attention( h, positions, step, @@ -265,7 +280,7 @@ def attention(self, h: torch.Tensor, positions: torch.Tensor, step: dict) -> tor head_dim=self.head_dim, eps=self.eps, rope_theta=self.rope_theta, - window=self.window, + window=self.window if self.is_sliding else step["max_seq_len"], ) @torch.no_grad() 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/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 index 2ed904ecba02..0ee2f234065a 100644 --- 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 @@ -72,7 +72,7 @@ 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 _sliding_attention +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, ) @@ -559,7 +559,7 @@ def check(name: str, rule: str | None, actual: torch.Tensor, expected: str) -> d head_size=d, attention_window_size=BLOCK_CONFIG["sliding_window"], **step, - **_sliding_attention.CALL_CONSTANTS, + **_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")) @@ -884,10 +884,9 @@ def test_rejects_unsupported_configuration(match, overrides) -> None: def test_rejects_other_blocks_and_devices() -> None: - with pytest.raises(ValueError, match="only block 0"): - BeamDenseBlockBf16Sm103Tp1( - block_config(layer_types=["sliding_attention"] * 2), 1, device="cuda" - ) + 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") 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 index 5b074b1bb50a..117bbe55508b 100644 --- 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 @@ -759,7 +759,7 @@ def refuse(self, *args, **kwargs): ("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 block 3"), + ("block_2", {}, 2, "only the routers of sparse blocks"), ] 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 index 1ff476a6ae63..be5e7f2c7fdb 100644 --- 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 @@ -449,7 +449,7 @@ def test_construction_allocates_only_the_router_state() -> None: (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 block 3"), + ({}, 0, "only the routers of sparse blocks"), ], ) def test_construction_rejects_unsupported_configs(overrides, layer_idx, match) -> None: 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_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..62c30777f257 --- /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,777 @@ +# 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 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 = {} + + +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 == max(LENGTHS) + 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, 2)[:, :130] + positions = positions_of(130) + 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/test_sparse_block.py b/tests/unittest/_torch/modeling_v2/models/beam/sparse_block_bf16__sm_103__tp1/test_sparse_block.py index 6e149331fe18..14beebaabcfc 100644 --- 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 @@ -982,8 +982,8 @@ def test_rejects_unsupported_configuration(case, overrides, match) -> None: def test_rejects_other_blocks_and_devices() -> None: - for layer_idx in (0, 5): - with pytest.raises(ValueError, match="only block 3"): + 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") @@ -1086,7 +1086,6 @@ def test_rejected_load_changes_nothing(case, error, match) -> None: "kv_cache_heads", "kv_cache_dtype", "layer0_cache", - "two_layer_cache", ] @@ -1094,7 +1093,7 @@ def test_rejected_load_changes_nothing(case, error, match) -> None: 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) and one that holds layer 3 beside another layer.""" + one-layer manager holds layer 0).""" blk = loaded_block(17) length = 17 x = torch.randn(1, length, 5120, device="cuda").to(DTYPE) @@ -1136,17 +1135,11 @@ def test_rejects_unsupported_inputs_before_dispatch(case) -> None: context = dict(num_tokens=length, dtype=DataType.HALF) elif case == "layer0_cache": context = dict(num_tokens=length, layer_mask=None) - elif case == "two_layer_cache": - context = dict( - num_tokens=length, num_layers=2, layer_mask=[i in (3, 4) for i in range(NUM_LAYERS)] - ) 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} - elif case == "two_layer_cache": - assert dict(md.kv_cache_manager.layer_offsets) == {3: 0, 4: 1} with no_dispatch(), pytest.raises(ValueError): blk(x, positions, md) @@ -1216,6 +1209,32 @@ def test_cache_slot_is_the_managers_local_index(ieee_fp32) -> None: 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 From 603bbb199e807e3df8dbe70cfafa9a4630dca61a Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Wed, 21 Oct 2026 10:01:58 -0700 Subject: [PATCH 08/13] [None][feat] Add cached generation and prefix reuse for Beam Drive the six-layer ModelingV2 component with KVCacheManagerV2 and TorchSampler for eager single-GPU greedy generation and shared-prefix reuse. Validate request history and model-weight identity before cache access, and add generated-data lifecycle and cache-safety tests. Signed-off-by: Maxim Gluhovskoi --- .../catalog/attention/thop_attention.py | 64 ++ .../modeling_v2/models/beam/_attention.py | 115 ++- .../dense_block_bf16__sm_103__tp1/block.py | 25 +- .../generation.py | 694 +++++++++++++++++ .../model.py | 261 +++++-- .../sparse_block_bf16__sm_103__tp1/block.py | 32 +- .../test_modeling_v2_thop_attention.py | 8 +- .../test_dense_block.py | 2 +- .../test_cached_generation.py | 723 ++++++++++++++++++ .../test_small_model_prefill.py | 6 +- .../test_sparse_block.py | 2 +- 11 files changed, 1829 insertions(+), 103 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/generation.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_cached_generation.py 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 fa28d4e05301..c319b9d1957c 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 @@ -85,6 +85,19 @@ 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) + #: 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( @@ -434,6 +447,52 @@ class ThopAttention(OpWrapper): ) 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 + ), + ), ) note = """ @@ -454,6 +513,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/models/beam/_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py index 8e25b9519916..ae8f77edb1ab 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py @@ -1,6 +1,6 @@ # 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 fresh context. +"""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 @@ -19,6 +19,17 @@ 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. +One call is one 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. + 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 @@ -29,8 +40,9 @@ 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 position check compares the positions with ``arange(T)`` on the -device and reads the answer back, once per call. The metadata checks read host tensors only. +Host synchronization: the position check 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 @@ -163,10 +175,11 @@ def check_step( """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(T)`` - on that device; a prepared ``TrtllmAttentionMetadata`` holding one fresh context - request of T tokens on a BF16 ``KVCacheManagerV2`` that holds logical layer - ``layer_idx`` with ``num_kv_heads`` kv heads of ``head_dim``. + 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") @@ -182,7 +195,7 @@ def check_step( 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 context lengths are 1..{max_tokens}, got {num_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() @@ -192,12 +205,17 @@ def check_step( raise ValueError( f"position_ids must be an integer [1, {num_tokens}] tensor on the block's device" ) - expected = torch.arange(num_tokens, dtype=position_ids.dtype, device=position_ids.device) - if not torch.equal(position_ids[0], expected): - raise ValueError("a fresh context must have positions 0..T-1") - check_metadata( + 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( @@ -207,33 +225,73 @@ def check_metadata( layer_idx: int, num_kv_heads: int, head_dim: int, -) -> None: - """Raise ValueError unless ``md`` is one prepared fresh context request of ``num_tokens`` - tokens on a BF16 ``KVCacheManagerV2`` holding logical layer ``layer_idx`` with - ``num_kv_heads`` kv heads of ``head_dim``. Reads host-side state only. +) -> int: + """Raise ValueError unless ``md`` is one prepared step of one request; return ``past``. + + 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 or + speculative decoding, one sequence whose KV length L fits the manager's + ``max_seq_len``, and either + + * a context step of ``num_tokens`` new tokens: ``num_contexts == 1``, context length + ``num_tokens``, ``past = L - num_tokens >= 0`` tokens already cached, and + paged-context FMHA whenever ``past > 0``; or + * a generation step: ``num_contexts == 0``, one new token after ``past = L - 1 >= 1``. + + ``past`` is the absolute position of the step's first new token. Reads host-side + state only. """ if not isinstance(md, TrtllmAttentionMetadata): raise ValueError(f"attn_metadata must be TrtllmAttentionMetadata, got {type(md).__name__}") mgr = md.kv_cache_manager if mgr is None: raise ValueError("attn_metadata needs a KV cache manager for the paged append") - if md.num_seqs != 1 or md.num_contexts != 1 or md.num_ctx_tokens != num_tokens: + if md.is_cross or md.is_spec_decoding_enabled or md.beam_width != 1: + raise ValueError("cross attention, speculative decoding and beams are not supported") + if md.num_seqs != 1: + raise ValueError(f"one request per call is supported; metadata has {md.num_seqs}") + if md.kv_cache_params is None or md.kv_cache_params.num_extra_kv_tokens: + raise ValueError("the metadata must use the KV cache without extra KV tokens") + if num_tokens > md.max_num_tokens: raise ValueError( - f"one context request of {num_tokens} tokens is supported; metadata has " - f"{md.num_seqs} sequences, {md.num_contexts} contexts, {md.num_ctx_tokens} context tokens" + f"the call has {num_tokens} tokens; the metadata holds {md.max_num_tokens} per call" ) - if int(md.kv_lens_runtime[0]) != num_tokens or int(md.prompt_lens_cpu_runtime[0]) != num_tokens: + kv_len = int(md.kv_lens_runtime[0]) + request_type = int(md.host_request_types_runtime[0]) + past = kv_len - num_tokens + if md.num_contexts == 1: + if ( + md.num_ctx_tokens != num_tokens + or int(md.prompt_lens_cpu_runtime[0]) != num_tokens + or request_type != 0 + ): + raise ValueError( + f"a context step of {num_tokens} tokens needs {num_tokens} context tokens and " + f"context length {num_tokens}; metadata has {md.num_ctx_tokens} context tokens " + f"and context length {int(md.prompt_lens_cpu_runtime[0])}" + ) + if past < 0: + raise ValueError( + f"the KV length {kv_len} is shorter than the call's {num_tokens} 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" + ) + elif md.num_contexts == 0: + if num_tokens != 1 or md.num_ctx_tokens != 0 or request_type != 1: + raise ValueError(f"a generation step appends exactly one token, got {num_tokens}") + if past < 1: + raise ValueError("a generation step needs the request's earlier tokens in the cache") + else: raise ValueError( - "a cached prefix is not supported: the request must start at an empty cache" + f"one request per call is supported; metadata has {md.num_contexts} contexts" ) - if ( - md.use_paged_context_fmha - or md.is_cross - or md.is_spec_decoding_enabled - or md.beam_width != 1 - ): + if kv_len > mgr.max_seq_len or md.max_seq_len != mgr.max_seq_len: raise ValueError( - "paged-context FMHA, cross attention, speculative decoding and beams are not supported" + f"the request reaches {kv_len} tokens; the cache holds {mgr.max_seq_len} per sequence " + f"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 [] @@ -243,6 +301,7 @@ def check_metadata( 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 past def step_args(md: TrtllmAttentionMetadata, layer_idx: int) -> dict: 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 index 92b1ba87658e..5ef87a5efefb 100644 --- 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 @@ -5,8 +5,9 @@ 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 fresh BF16 context through it, -with the attention batch state prepared on a ``KVCacheManagerV2``. The forward is +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. @@ -68,8 +69,9 @@ SUPPORTED_LAYER_INDICES = (0, 1, 2) SUPPORTED_ROPE = {"rope_type": "default", "rope_theta": 10000.0} -#: The longest fresh context the block was measured at; longer ones are rejected. -MAX_TOKENS = 129 +#: 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) @@ -146,15 +148,16 @@ def check_config(config, layer_idx: int) -> None: class BeamDenseBlockBf16Sm103Tp1(nn.Module): - """Beam dense decoder block 0, 1 or 2: BF16 weights and activations, one GPU, one fresh context. + """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 <= 129, positions ``arange(T)`` and a - prepared ``TrtllmAttentionMetadata`` holding exactly that one context request on a - BF16 ``KVCacheManagerV2`` (8 kv heads of 128, no cached prefix) that holds the - block's logical layer ``layer_idx``. Inputs outside that are rejected before any - kernel runs; decode, cached prefixes, chunked prefill, batches and other blocks are - not computed here. + 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__( 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 index e6f527d6e382..9b7ecaa2daaa 100644 --- 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 @@ -1,14 +1,16 @@ # 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 prefill logits, constructed directly. +"""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. 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 fresh BF16 context through it, -with the attention batch state prepared on a six-layer ``KVCacheManagerV2`` -(``prefill_context`` builds one). It does not call the built-in Beam modeling. +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:: @@ -23,28 +25,42 @@ 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 keep a 512-token window and layer 5 keeps the whole context; the -manager states this as its per-layer window vector ``[512] * 5 + [None]`` (``None`` is +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. ``check_cache`` verifies all of it before the first layer runs, including that -the host pool tables the attention op reads address each slot's own key and value pages -as the manager allocated them. Cache contents are written as the backend requires; they -are not read back, reused or shared between calls. - -Host synchronization per call: one for the model's token-id range and position 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. +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 @@ -59,7 +75,7 @@ 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 LayerId, PageIndexMode +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 @@ -86,14 +102,15 @@ "final_logit_softcapping": 30.0, } -#: The longest fresh context the model was measured at; longer ones are rejected. -MAX_TOKENS = 129 - #: 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'. @@ -149,16 +166,17 @@ def weight_schema(config) -> "OrderedDict[str, tuple[tuple[int, ...], torch.dtyp return schema -def check_cache(md: TrtllmAttentionMetadata, num_tokens: int) -> None: +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 call's tokens; keys and values in one shared page - table per pool; and 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``). Reads host-side manager state and host tensors only. + 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): @@ -199,12 +217,113 @@ def check_cache(md: TrtllmAttentionMetadata, num_tokens: int) -> None: 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 < num_tokens or md.max_seq_len != mgr.max_seq_len: + 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 call has {num_tokens}" + 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: @@ -354,15 +473,22 @@ class BeamSmallModelPrefillBf16Sm103Tp1(nn.Module): 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 token t's - softcapped logit for vocabulary id v. Accepted: 1 <= T <= 129 token ids in - ``[0, 151936)``, contiguous on the model's device; integer positions equal to - ``arange(T)``; and the prepared ``TrtllmAttentionMetadata`` of one fresh context - request of T tokens on the six-layer cache ``check_cache`` describes, such as - ``prefill_context(T)`` yields. 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. Decode, cached prefixes, chunked prefill, - batches, padding and contexts longer than 129 tokens are not computed here. + ``[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"): @@ -393,6 +519,8 @@ def empty(name: str) -> nn.Parameter: 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: @@ -427,7 +555,17 @@ def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: 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)): @@ -439,6 +577,14 @@ def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: ) 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: @@ -455,7 +601,9 @@ def _check_call( 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 context lengths are 1..{MAX_TOKENS}, got {num_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 ( @@ -472,18 +620,45 @@ def _check_call( raise ValueError( f"attn_metadata must be TrtllmAttentionMetadata, got {type(md).__name__}" ) - check_cache(md, num_tokens) - for layer_idx in range(NUM_LAYERS): + 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 ) - expected = torch.arange(num_tokens, dtype=position_ids.dtype, device=self.device) + 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 = torch.stack( - ((position_ids[0] == expected).all(), (low >= 0) & (high < self.vocab_size)) + 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("a fresh context must have positions 0..T-1") + 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 " 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 index 325f763c51d1..7746f948c39b 100644 --- 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 @@ -5,9 +5,10 @@ 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 fresh BF16 context through it, 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 +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. @@ -37,8 +38,9 @@ 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). Lengths up to 129 tokens never -reach the 512-token window, so behaviour at or beyond the window is not exercised here. +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 @@ -82,8 +84,9 @@ #: RoPE of the sliding layers; the full-attention layers have none. SUPPORTED_ROPE = {"rope_type": "default", "rope_theta": 10000.0} -#: The longest fresh context the block was measured at; longer ones are rejected. -MAX_TOKENS = 129 +#: 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) @@ -159,18 +162,19 @@ def check_config(config, layer_idx: int) -> None: class BeamSparseBlockBf16Sm103Tp1(nn.Module): """Beam sparse decoder block 3, 4 or 5: BF16 weights and activations, an FP32 router, one GPU, - one fresh context. + 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 <= 129, positions ``arange(T)`` and a - prepared ``TrtllmAttentionMetadata`` holding exactly that one context request on a - BF16 ``KVCacheManagerV2`` (8 kv heads of 128, no cached prefix) that holds the - block's logical layer. Inputs outside that are rejected before any kernel runs; - decode, cached prefixes, chunked prefill, batches and other blocks are not computed - here. Inputs and parameters are never modified. + 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__( 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 7a3d08f2af83..95475e710b32 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 @@ -5390,6 +5390,7 @@ def _op_reference( causal: bool = True, e4m3: bool = False, layer: Optional[int] = None, + window: Optional[int] = None, ) -> torch.Tensor: """`op.reference` over the history the env mirrored for these requests. @@ -5414,6 +5415,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, ) @@ -6944,10 +6946,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 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"] @@ -6984,7 +6987,8 @@ def _cell_paged(cell, arm) -> None: assert not (fp8 and not paged and 0 < nc < len(rids)) e4m3 = fp8 and (paged or nc == 0) causal = spec["mask_type"] == MASK_CAUSAL - _compare(out, _op_reference(env, qkv, list(new), rids, causal=causal, e4m3=e4m3)) + ref = _op_reference(env, qkv, list(new), rids, causal=causal, e4m3=e4m3, window=spec["window"]) + _compare(out, ref) for rid in rids: env.check_cache(rid) env.check_unwritten_pool_zero() 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 index 0ee2f234065a..618859de29d2 100644 --- 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 @@ -797,7 +797,7 @@ def test_configuration_is_the_fixed_scope() -> None: 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 == max(LENGTHS) + 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) == ( 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..f8a621534b95 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/small_model_prefill_bf16__sm_103__tp1/test_cached_generation.py @@ -0,0 +1,723 @@ +# 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 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 = {} + + +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 index 62c30777f257..8b0eb5f3251c 100644 --- 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 @@ -368,7 +368,7 @@ def test_configuration_is_the_fixed_scope() -> None: 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 == max(LENGTHS) + 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()) @@ -573,8 +573,8 @@ def test_rejects_unsupported_inputs_before_dispatch(case, monkeypatch) -> None: elif case == "empty": ids, positions = ids[:, :0], positions[:, :0] elif case == "too_long": - ids = generate_ids(17, 129).repeat(1, 2)[:, :130] - positions = positions_of(130) + 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": 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 index 14beebaabcfc..b1d0fa65ae40 100644 --- 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 @@ -899,7 +899,7 @@ def test_configuration_is_the_fixed_scope() -> None: 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 == max(LENGTHS) == reference.MAX_TOKENS + 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) From 55f78a52da72216510c4be7e31a3844705cd56cd Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Wed, 21 Oct 2026 17:10:14 -0700 Subject: [PATCH 09/13] [None][feat] Add Beam checkpoint loading and TP4 block execution Add bounded streaming checkpoint loading, rank-local ownership, and representative distributed block execution on four GPUs. Include generated-data loader/layout tests and a model-specific RMSNorm operator. This qualifies checkpoint-derived BF16 execution and full parameter residency; full-depth inference and native quantized execution remain outside this component. Signed-off-by: Maxim Gluhovskoi --- .../__init__.py | 3 + .../layout.py | 369 ++++++++++ .../loader.py | 338 +++++++++ .../model.py | 665 ++++++++++++++++++ .../ops/__init__.py | 3 + .../ops/norm/__init__.py | 2 + .../ops/norm/beam_rmsnorm_fp32.py | 78 ++ .../test_modeling_v2_beam_rmsnorm_fp32.py | 76 ++ .../test_checkpoint_layout.py | 490 +++++++++++++ 9 files changed, 2024 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/layout.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/loader.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/model.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/norm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/ops/norm/beam_rmsnorm_fp32.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/norm/test_modeling_v2_beam_rmsnorm_fp32.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/test_checkpoint_layout.py 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..b6efb0d5bf48 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_bf16__sm_103__tp4/model.py @@ -0,0 +1,665 @@ +# 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 +#: Tokens per call; the representative cases use at most 129. +MAX_TOKENS = 1024 + + +class LoadError(RuntimeError): + """A load failed on at least one rank; the component is not usable on any rank.""" + + +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). 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, + ): + super().__init__() + 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 (OSError, ValueError) as exc: + 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 = TileReader() + try: + 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 (OSError, ValueError, RuntimeError) as exc: + error = exc + finally: + bytes_read, reads = reader.bytes_read, reader.reads + reader.close() + 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) + 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, + ) -> torch.Tensor: + """Normalized input ``[T, hidden]`` -> attention output after the one o_proj all-reduce.""" + self._check_usable(layer) + _check_hidden(h, self.hidden, self.device) + T = h.shape[0] + _attention.check_metadata( + md, T, layer_idx=layer, num_kv_heads=self.local_kv_heads, head_dim=self.head_dim + ) + 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, + ) + 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) + 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) + 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, + ) -> torch.Tensor: + """One decoder block of logical layer ``layer`` on ``x`` ``[T, hidden]``, one request 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. + """ + 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) + 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) + 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) + 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) -> 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) -> 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("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", + "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/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) From 04b623524404d5be6a16ef0e4e94c5e18ca2ec18 Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Thu, 22 Oct 2026 09:31:08 -0700 Subject: [PATCH 10/13] [None][feat] Add Beam NVFP4 expert loading and TP4 execution Load packed NVFP4 expert weights and scales into rank-local storage and execute representative sparse blocks with TRT-LLM CUTLASS kernels. Preserve BF16 excluded modules and FP32 routers, and add generated-data format, scale, sharding and rejection tests. Qualifies full 52-layer parameter residency and representative eager component execution on four GB300 GPUs. Full-depth inference, generation, task accuracy, CUDA graphs and performance qualification remain future work. Signed-off-by: Maxim Gluhovskoi --- .../__init__.py | 3 + .../expert.py | 129 ++++ .../layout.py | 335 ++++++++++ .../loader.py | 170 +++++ .../model.py | 485 ++++++++++++++ .../test_checkpoint_layout_nvfp4.py | 624 ++++++++++++++++++ 6 files changed, 1746 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/expert.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/layout.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/loader.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/model.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/test_checkpoint_layout_nvfp4.py 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..10882677d780 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/checkpoint_layout_nvfp4__sm_103__tp4/model.py @@ -0,0 +1,485 @@ +# 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 ( + BeamCheckpointLayoutBf16Sm103Tp4, + LoadError, + _attr, + _check_mapping, + _rec, +) +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 (OSError, ValueError) as exc: + 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 = TileReader() + + def first_entry(spec): + return catalog.entries[ + spec.name + ".weight" if isinstance(spec, QuantSpec) else spec.name + ] + + 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)] + try: + 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 (OSError, ValueError, RuntimeError) as exc: + error = exc + finally: + bytes_read, reads = reader.bytes_read, reader.reads + reader.close() + 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/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) From bfc864a7f69cfa7e8b0d1c593ec701c20e1b367b Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Fri, 23 Oct 2026 13:05:06 -0700 Subject: [PATCH 11/13] [None][feat] Add Beam full-model NVFP4 generation and serving Integrate all 52 layers with the ModelingV2 target, native checkpoint loading, KV cache management, TorchSampler and text chat-completions serving. Support eager TP4 execution and native request batching; add parser/routing and batch metadata coverage, a launch guide and CPU CI registration. Complete the registered 10-shot GSM8K evaluation: 1263/1319 (95.75%), compared with a recorded BF16 mean of 95.53%. Batching numerical qualification remains partial; serial-equivalent trajectories and optimized throughput are not claimed. Independent reviewer and QA accepted the scoped functional baseline. Signed-off-by: Maxim Gluhovskoi --- .../modeling_v2/_router_index.py | 1 + .../modeling_v2/models/beam/README.md | 116 ++++++++ .../modeling_v2/models/beam/_attention.py | 144 +++++---- .../model.py | 59 ++-- .../model.py | 10 +- .../full_model_nvfp4__sm_103__tp4/__init__.py | 2 + .../full_model_nvfp4__sm_103__tp4/model.py | 278 ++++++++++++++++++ .../full_model_nvfp4__sm_103__tp4/modeling.py | 185 ++++++++++++ .../full_model_nvfp4__sm_103__tp4/serving.py | 264 +++++++++++++++++ .../full_model_nvfp4__sm_103__tp4/weights.py | 52 ++++ .../modeling_v2/models/beam/routing.py | 61 ++++ .../_torch/pyexecutor/config_utils.py | 7 + .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../test_batching.py | 144 +++++++++ .../test_serving.py | 134 +++++++++ 15 files changed, 1387 insertions(+), 71 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/model.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/serving.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/weights.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/routing.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_batching.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_serving.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index b0b6e78ef8d8..84771dad522f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -49,6 +49,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/models/beam/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md new file mode 100644 index 000000000000..2e21a6f6cdfa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md @@ -0,0 +1,116 @@ +# 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`; greedy decoding is the evaluated policy | +| 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. + +## 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. + +## 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, speculative decoding, +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/_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py index ae8f77edb1ab..bfe5d3a28afb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py @@ -19,9 +19,9 @@ 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. -One call is one 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 +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 @@ -30,6 +30,11 @@ 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. + 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 @@ -40,9 +45,9 @@ 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 position check 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. +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 @@ -228,18 +233,45 @@ def check_metadata( ) -> int: """Raise ValueError unless ``md`` is one prepared step of one request; return ``past``. - Accepted on a BF16 ``KVCacheManagerV2`` holding logical layer ``layer_idx`` with + ``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, +) -> 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 or - speculative decoding, one sequence whose KV length L fits the manager's - ``max_seq_len``, and either + speculative decoding, where each request ``i`` with KV length ``L_i`` (at most the + manager's ``max_seq_len``) is - * a context step of ``num_tokens`` new tokens: ``num_contexts == 1``, context length - ``num_tokens``, ``past = L - num_tokens >= 0`` tokens already cached, and - paged-context FMHA whenever ``past > 0``; or - * a generation step: ``num_contexts == 0``, one new token after ``past = L - 1 >= 1``. + * 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: one new token after ``past_i = L_i - 1 >= 0`` cached tokens + (the native runtime's warmup issues generation steps over a cache that holds only the + new token). - ``past`` is the absolute position of the step's first new token. Reads host-side - state only. + 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__}") @@ -248,50 +280,60 @@ def check_metadata( 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, speculative decoding and beams are not supported") - if md.num_seqs != 1: - raise ValueError(f"one request per call is supported; metadata has {md.num_seqs}") + 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 or md.kv_cache_params.num_extra_kv_tokens: raise ValueError("the metadata must use the KV cache without extra KV tokens") 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" ) - kv_len = int(md.kv_lens_runtime[0]) - request_type = int(md.host_request_types_runtime[0]) - past = kv_len - num_tokens - if md.num_contexts == 1: - if ( - md.num_ctx_tokens != num_tokens - or int(md.prompt_lens_cpu_runtime[0]) != num_tokens - or request_type != 0 - ): - raise ValueError( - f"a context step of {num_tokens} tokens needs {num_tokens} context tokens and " - f"context length {num_tokens}; metadata has {md.num_ctx_tokens} context tokens " - f"and context length {int(md.prompt_lens_cpu_runtime[0])}" - ) - if past < 0: - raise ValueError( - f"the KV length {kv_len} is shorter than the call's {num_tokens} 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" - ) - elif md.num_contexts == 0: - if num_tokens != 1 or md.num_ctx_tokens != 0 or request_type != 1: - raise ValueError(f"a generation step appends exactly one token, got {num_tokens}") - if past < 1: - raise ValueError("a generation step needs the request's earlier tokens in the cache") - else: + 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"one request per call is supported; metadata has {md.num_contexts} contexts" + f"the requests' new tokens {new_tokens} ({md.num_ctx_tokens} in contexts) do not " + f"fill the call's {num_tokens} tokens" ) - if kv_len > mgr.max_seq_len or md.max_seq_len != mgr.max_seq_len: + 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 != 1 or request_type != 1: + raise ValueError( + f"generation request {i} appends exactly one token, got {n} (request type " + f"{request_type})" + ) + if past < 0: + raise ValueError( + f"request {i}'s KV length {kv_len} is shorter than its generation token" + ) + spans.append((past, n)) + if max(kv_lens) > mgr.max_seq_len or md.max_seq_len != mgr.max_seq_len: raise ValueError( - f"the request reaches {kv_len} tokens; the cache holds {mgr.max_seq_len} per sequence " - f"and the metadata states {md.max_seq_len}" + 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 [] @@ -301,7 +343,7 @@ def check_metadata( 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 past + return spans def step_args(md: TrtllmAttentionMetadata, layer_idx: int) -> dict: 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 index b6efb0d5bf48..2fe059ddd875 100644 --- 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 @@ -92,7 +92,8 @@ _SM = (10, 3) ROUTER_EPS = 1e-20 -#: Tokens per call; the representative cases use at most 129. +#: 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 @@ -100,6 +101,13 @@ 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) @@ -131,8 +139,9 @@ class BeamCheckpointLayoutBf16Sm103Tp4(nn.Module): ``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). Construction allocates uninitialized rank-local storage on the current - CUDA device and validates the configuration and layout; it reads no checkpoint data. + 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__( @@ -141,8 +150,11 @@ def __init__( 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 = ( @@ -360,7 +372,7 @@ def _digests_agree(self, tensors: Sequence[torch.Tensor]) -> list: 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) + _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( @@ -380,11 +392,16 @@ def attention( md: TrtllmAttentionMetadata, capture: Optional[dict] = None, ) -> torch.Tensor: - """Normalized input ``[T, hidden]`` -> attention output after the one o_proj all-reduce.""" + """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``) + 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) + _check_hidden(h, self.hidden, self.device, self.max_tokens) T = h.shape[0] - _attention.check_metadata( + _attention.check_batch_metadata( md, T, layer_idx=layer, num_kv_heads=self.local_kv_heads, head_dim=self.head_dim ) step = _attention.step_args(md, layer) @@ -466,7 +483,7 @@ def dense_ffn( ) -> 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) + _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()) @@ -492,7 +509,7 @@ def dense_ffn( 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) + _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( @@ -569,7 +586,7 @@ def block( md: TrtllmAttentionMetadata, capture: Optional[dict] = None, ) -> torch.Tensor: - """One decoder block of logical layer ``layer`` on ``x`` ``[T, hidden]``, one request step. + """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. @@ -606,7 +623,7 @@ def block( @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) + _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 @@ -619,7 +636,7 @@ def head(self, n: torch.Tensor, capture: Optional[dict] = None) -> torch.Tensor: 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) + _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) @@ -638,22 +655,26 @@ def _rec(capture: Optional[dict], items: dict) -> None: capture[name] = tensor.detach().clone() -def _check_hidden(h: torch.Tensor, hidden: int, device: torch.device) -> None: +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: + 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)}" + 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) -> None: +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("token ids must hold 1..MAX_TOKENS ids on the component's device") + 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})") @@ -661,5 +682,7 @@ def _check_ids(ids: torch.Tensor, vocab: int, device: torch.device) -> None: __all__ = [ "BeamCheckpointLayoutBf16Sm103Tp4", "LoadError", + "MAX_TOKENS", + "check_max_tokens", "TensorSpec", ] 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 index 10882677d780..4002a1fd96cf 100644 --- 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 @@ -50,11 +50,13 @@ 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 @@ -99,8 +101,9 @@ class BeamCheckpointLayoutNvfp4Sm103Tp4(BeamCheckpointLayoutBf16Sm103Tp4): ``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). Construction allocates uninitialized rank-local storage on the current CUDA device - and validates the configuration and layout; it reads no checkpoint data. + 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__( @@ -109,8 +112,11 @@ def __init__( mapping: ParallelMapping, layers: Optional[Sequence[int]] = None, device: Optional[torch.device] = None, + *, + max_tokens: int = MAX_TOKENS, ): nn.Module.__init__(self) + self.max_tokens = check_max_tokens(max_tokens) self.schedule = check_config(config) _check_mapping(mapping) device = ( 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..3edfe73d465d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/model.py @@ -0,0 +1,278 @@ +# 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 one token. 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, or explicit row +indices. + +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; 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]) -> 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": + raise ValueError( + f"logit_rows must be None, 'all' or row indices, got {logit_rows!r}" + ) + return list(range(num_tokens)) + 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, + ) -> 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``; ``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, + ) + 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, + ) -> torch.Tensor: + """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). + """ + self._check_usable() + ids, positions, spans = self.check_call(input_ids, position_ids, attn_metadata) + num_tokens = ids.shape[0] + rows = self._select_rows(logit_rows, spans) + 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) + 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 + + +__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..79afd7af669a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/modeling.py @@ -0,0 +1,185 @@ +# 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, no speculative decoding) when ``modeling_v2`` routes it. It owns exactly one +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. +""" + +from __future__ import annotations + +import json +import os +from typing import Optional + +import torch +from torch import nn + +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" + ) + + +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) -> torch.Tensor: + kind = "batched" if attn_metadata.num_seqs > 1 else "single-request" + trace = None if kind in self.traced else [] + logits = self.model( + input_ids, position_ids, attn_metadata, logit_rows=logit_rows, trace=trace + ) + if trace is not None: + self.traced.add(kind) + _log_trace(f"{kind} {self.phase}", trace) + return logits + + +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 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 + 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" + ) + + @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 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, + **kwargs, + ) -> torch.Tensor: + 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("spec_metadata") is not None or kwargs.get("lora_params") is not None: + raise ValueError("speculative decoding and LoRA are not supported by this target") + ids = input_ids.reshape(-1).to(torch.int64) + 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 + ) + + +__all__ = ["DecodeTarget", "ModelingV2BeamFullModelNvfp4Sm103Tp4", "PrefillTarget"] 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..59ac84038609 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/serving.py @@ -0,0 +1,264 @@ +# 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 [--port 8000] [--max-seq-len 69632] [--max-num-tokens 8192] \\ + [--max-batch-size 1] [--iter-stats] + +``--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 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, +) -> 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. + """ + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + + from .weights import BeamNvfp4CheckpointLoader + + return { + "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, + } + + +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") + 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, + ), + 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..de0bbfd9a96a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/weights.py @@ -0,0 +1,52 @@ +# 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. +""" + +from __future__ import annotations + +from typing import Any + +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 from its checkpoint directory.""" + + def __init__(self): + super().__init__() + self._preloaded = False + + 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 not isinstance(model, ModelingV2BeamFullModelNvfp4Sm103Tp4): + raise TypeError( + "BeamNvfp4CheckpointLoader loads only ModelingV2BeamFullModelNvfp4Sm103Tp4 " + f"(modeling_v2='require'); got {type(model).__name__}" + ) + self._preloaded = False + model.load_checkpoint(checkpoint_dir) + self._preloaded = True + return {} + + def is_weights_preloaded(self) -> bool: + return self._preloaded + + +__all__ = ["BeamNvfp4CheckpointLoader"] 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..266220416826 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/routing.py @@ -0,0 +1,61 @@ +# 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 Any, Optional, Tuple + +from ..._router_index import NoTarget + + +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 route( + pretrained_config: Any, + mapping: Any, + compute_capability: Tuple[int, int], + spec_config: Optional[Any], +) -> 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, no attention data parallelism, no speculative decoding. + """ + 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) + if ( + compute_capability == (10, 3) + and shape == (52, 5120, 256) + and layout == (4, 4, 1, 4, 1, False) + and quant_algo == "NVFP4" + and spec_config 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'}; " + "the one Beam target is the 52-layer NVFP4 checkpoint on sm103, attention TP4, MoE " + "TP4/EP1, PP1, without attention DP or speculative decoding" + ) diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index a9b2b16bfbe3..ffef9828edbc 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -1118,6 +1118,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_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 58ebb7ca4eba..9d0835b5c607 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -59,6 +59,7 @@ l0_cpu: # compute capability, no device. 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/models/test_minimax_m3.py::test_minimax_m3_dense_backend_library_order - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_moe_reduces_only_local_terms 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..1e433991a4e2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_batching.py @@ -0,0 +1,144 @@ +# 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.""" + 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_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..83a237f11363 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_serving.py @@ -0,0 +1,134 @@ +# 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).""" + +from types import SimpleNamespace + +import pytest + +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( + 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) + assert target.__name__ == "ModelingV2BeamFullModelNvfp4Sm103Tp4" + declined = [ + (_beam_config(), _mapping(), (10, 0), None), + (_beam_config(num_hidden_layers=6), _mapping(), (10, 3), None), + (_beam_config(quantization_config=None), _mapping(), (10, 3), None), + (_beam_config(), _mapping(tp_size=2, pp_size=2, moe_tp_size=2), (10, 3), None), + (_beam_config(), _mapping(moe_tp_size=1, moe_ep_size=4), (10, 3), None), + (_beam_config(), _mapping(enable_attention_dp=True), (10, 3), None), + (_beam_config(), _mapping(), (10, 3), object()), + ] + for args in declined: + with pytest.raises(NoTarget): + route(*args) From 96ce0eca1be9f3134ee3fa3a4e3e9d99a0709ca9 Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Fri, 23 Oct 2026 18:05:09 -0700 Subject: [PATCH 12/13] [None][fix] Harden Beam loading and serving validation Keep rank-local metadata, streaming and cleanup failures inside the all-rank load agreement. Validate dense-block weights before copying and reject execution until a complete load succeeds. Reject unsupported launcher token budgets and tensor parallelism in the unsharded loader. Make the fallback-selection test independent of the host GPU and skip CUDA metadata setup when no device is available. Add CPU regression coverage for these contracts and Beam config loading. Validation: 31 focused CPU cases passed with native imports isolated; repository pre-commit checks passed. Native GPU execution was not rerun. Signed-off-by: Maxim Gluhovskoi --- .../model.py | 17 +- .../model.py | 25 +- .../dense_block_bf16__sm_103__tp1/block.py | 14 +- .../full_model_nvfp4__sm_103__tp4/serving.py | 3 + tensorrt_llm/_torch/models/modeling_beam.py | 5 + .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../test_flashinfer_trtllm_gen_fmha.py | 2 + .../test_batching.py | 2 + .../test_serving.py | 51 ++++ .../models/beam/test_loading_contracts.py | 219 ++++++++++++++++++ 10 files changed, 324 insertions(+), 15 deletions(-) create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/test_loading_contracts.py 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 index 2fe059ddd875..cd2068ce4c04 100644 --- 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 @@ -278,7 +278,7 @@ def load(self, checkpoint_dir: str, *, record_digests: bool = True) -> dict: ) catalog = SafetensorsCatalog(checkpoint_dir) catalog.validate(tensor_specs(self.config), complete=True) - except (OSError, ValueError) as exc: + except Exception as exc: # Keep rank-local failures inside the collective protocol. error = exc if not self._agree(error is None): raise LoadError( @@ -288,8 +288,10 @@ def load(self, checkpoint_dir: str, *, record_digests: bool = True) -> dict: self.usable = False self.load_record = None coverage = OrderedDict() - reader = TileReader() + 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), @@ -306,11 +308,16 @@ def load(self, checkpoint_dir: str, *, record_digests: bool = True) -> dict: "bank": spec.bank, "slot": spec.expert, } - except (OSError, ValueError, RuntimeError) as exc: + except Exception as exc: # Every rank must reach the completion agreement. error = exc finally: - bytes_read, reads = reader.bytes_read, reader.reads - reader.close() + 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( 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 index 4002a1fd96cf..5ee2a9e3ba1a 100644 --- 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 @@ -326,7 +326,7 @@ def load(self, checkpoint_dir: str, *, record_digests: bool = True) -> dict: recipe = check_recipe(quant, full) catalog = SafetensorsCatalog(checkpoint_dir) validate(catalog, stored_entries(full), complete=True) - except (OSError, ValueError) as exc: + except Exception as exc: # Keep rank-local failures inside the collective protocol. error = exc if not self._agree(error is None): raise LoadError( @@ -338,18 +338,20 @@ def load(self, checkpoint_dir: str, *, record_digests: bool = True) -> dict: self._shared_fc1 = {} coverage = OrderedDict() scalars = OrderedDict() - reader = TileReader() + reader = None + bytes_read = reads = 0 def first_entry(spec): return catalog.entries[ spec.name + ".weight" if isinstance(spec, QuantSpec) else spec.name ] - 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)] 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): @@ -365,11 +367,16 @@ def first_entry(spec): } for layer in sparse: scalars[layer] = self._load_scalars(catalog, reader, layer) - except (OSError, ValueError, RuntimeError) as exc: + except Exception as exc: # Every rank must reach the completion agreement. error = exc finally: - bytes_read, reads = reader.bytes_read, reader.reads - reader.close() + 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 ) 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 index 5ef87a5efefb..3900e91aaca8 100644 --- 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 @@ -183,6 +183,7 @@ def __init__( 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.""" @@ -191,17 +192,26 @@ def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": @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" ) - getattr(self, _attribute(name)).copy_(tensor) + 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 @@ -218,6 +228,8 @@ def _check_step( 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: 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 index 59ac84038609..fccde334b730 100644 --- 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 @@ -199,8 +199,11 @@ def llm_args( """ 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}") return { "model": checkpoint, "tokenizer": load_tokenizer(checkpoint), diff --git a/tensorrt_llm/_torch/models/modeling_beam.py b/tensorrt_llm/_torch/models/modeling_beam.py index a53c8294eb8b..1b41258d71d8 100644 --- a/tensorrt_llm/_torch/models/modeling_beam.py +++ b/tensorrt_llm/_torch/models/modeling_beam.py @@ -694,6 +694,11 @@ def load_weights( """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) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 5808d8c0178e..3b246cc4a0ff 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -67,6 +67,7 @@ l0_cpu: - 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 a32f2ce472f5..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,7 @@ 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 @@ -149,6 +150,7 @@ def kernel_exists(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 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 index 1e433991a4e2..83fdec764682 100644 --- 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 @@ -39,6 +39,8 @@ @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), 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 index cae9cb8ae836..80a59f75a72e 100644 --- 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 @@ -2,9 +2,11 @@ # 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 @@ -163,3 +165,52 @@ def test_engine_resolver_honors_effective_kv_format( 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 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, {})] From 57fed0768ec3c932174342e7378e73d1c5d24e86 Mon Sep 17 00:00:00 2001 From: Maxim Gluhovskoi Date: Sat, 24 Oct 2026 18:56:02 -0700 Subject: [PATCH 13/13] [None][feat] Add Beam DFlash2 speculative decoding Signed-off-by: Maxim Gluhovskoi --- .../catalog/attention/thop_attention.py | 56 +++ .../modeling_v2/models/beam/README.md | 47 +- .../modeling_v2/models/beam/_attention.py | 65 ++- .../model.py | 36 +- .../full_model_nvfp4__sm_103__tp4/model.py | 116 ++++- .../full_model_nvfp4__sm_103__tp4/modeling.py | 308 +++++++++++- .../full_model_nvfp4__sm_103__tp4/serving.py | 47 +- .../full_model_nvfp4__sm_103__tp4/weights.py | 31 +- .../modeling_v2/models/beam/routing.py | 56 ++- tensorrt_llm/_torch/models/modeling_dflash.py | 159 ++++++- .../test_lists/test-db/l0_b300.yml | 4 + .../test_modeling_v2_dflash_attention.py | 262 ++++++++++ .../test_modeling_v2_thop_attention.py | 147 ++++++ .../test_dflash2.py | 448 ++++++++++++++++++ .../test_serving.py | 143 ++++++ .../test_cached_generation.py | 10 + .../test_small_model_prefill.py | 10 + 17 files changed, 1856 insertions(+), 89 deletions(-) create mode 100644 tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_dflash_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/test_dflash2.py 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 6ac499776be3..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 @@ -127,6 +127,20 @@ def _spread(batch: int, longest: int) -> tuple[int, ...]: ) _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( @@ -853,6 +867,48 @@ class ThopAttention(OpWrapper): _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 = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md index 2e21a6f6cdfa..d1111389e62f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md @@ -15,7 +15,7 @@ ModelingV2 explicitly; it does not fall back to the built-in model implementatio | 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`; greedy decoding is the evaluated policy | +| 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 | @@ -64,6 +64,39 @@ native batching can be enabled with `--max-batch-size`. The benchmark used 128 a 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 @@ -96,6 +129,16 @@ 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 @@ -110,7 +153,7 @@ integration work; generated tests must not be presented as that coverage. 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, speculative decoding, +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/_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py index bfe5d3a28afb..ba71826dec2c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py @@ -35,6 +35,15 @@ 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 @@ -136,6 +145,16 @@ ) +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. @@ -253,40 +272,59 @@ def check_batch_metadata( 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 or - speculative decoding, where each request ``i`` with KV length ``L_i`` (at most the - manager's ``max_seq_len``) is + ``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: one new token after ``past_i = L_i - 1 >= 0`` cached tokens - (the native runtime's warmup issues generation steps over a cache that holds only the - new token). + * 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, speculative decoding and beams are not supported") + 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 or md.kv_cache_params.num_extra_kv_tokens: - raise ValueError("the metadata must use the KV cache without extra KV tokens") + 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" @@ -320,14 +358,15 @@ def check_batch_metadata( "paged-context FMHA is required" ) else: - if n != 1 or request_type != 1: + 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 one token, got {n} (request type " - f"{request_type})" + 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 generation token" + 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: 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 index cd2068ce4c04..b76bda2ca58c 100644 --- 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 @@ -398,18 +398,28 @@ def attention( 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``) - whose new tokens are the T rows, contexts first; ``positions`` are their absolute - positions. + ``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 + 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 @@ -439,7 +449,7 @@ def attention( head_size=d, attention_window_size=self.window if sliding else step["max_seq_len"], **step, - **_attention.CALL_CONSTANTS, + **_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") @@ -592,17 +602,29 @@ def block( 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. + 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) + 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) 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 index 3edfe73d465d..c2eeb529fb83 100644 --- 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 @@ -18,11 +18,16 @@ 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 one token. 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, or explicit row -indices. +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 @@ -31,8 +36,9 @@ 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; the -component's router agreement and expert dispatch synchronize per sparse layer. +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 @@ -129,7 +135,9 @@ def check_cache_manager(self, manager) -> None: f"{window} (max_seq_len {manager.max_seq_len})" ) - def _select_rows(self, logit_rows: LogitRows, spans: Sequence[tuple]) -> list: + 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 @@ -138,11 +146,18 @@ def _select_rows(self, logit_rows: LogitRows, spans: Sequence[tuple]) -> list: ends.append(end - 1) return ends if isinstance(logit_rows, str): - if logit_rows != "all": + if logit_rows == "all": + return list(range(num_tokens)) + if logit_rows != "verification": raise ValueError( - f"logit_rows must be None, 'all' or row indices, got {logit_rows!r}" + "logit_rows must be None, 'all', 'verification' or row indices, got " + f"{logit_rows!r}" ) - return list(range(num_tokens)) + 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 @@ -159,15 +174,19 @@ def check_call( 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``; ``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. + 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") @@ -205,6 +224,8 @@ def check_call( 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]) @@ -230,24 +251,52 @@ def forward( logit_rows: LogitRows = None, layer_capture: Optional[Callable[[int], Optional[dict]]] = None, trace: Optional[list] = None, - ) -> torch.Tensor: + 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). + 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) + 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) + 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) + 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( @@ -267,7 +316,32 @@ def forward( 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 + 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__ = [ 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 index 79afd7af669a..4d273c5583aa 100644 --- 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 @@ -4,11 +4,11 @@ ``ModelingV2BeamFullModelNvfp4Sm103Tp4`` is what ``AutoModelForCausalLM`` builds for a ``BeamForCausalLM`` NVFP4 checkpoint on four sm_103 ranks (attention TP4, MoE TP4/EP1, PP1, no -attention DP, no speculative decoding) when ``modeling_v2`` routes it. It owns exactly one -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. +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: @@ -19,21 +19,40 @@ * ``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. +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 @@ -68,6 +87,46 @@ def _log_trace(phase: str, trace: list) -> None: ) +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.""" @@ -77,16 +136,16 @@ def __init__(self, model: BeamFullModelNvfp4Sm103Tp4): self.model = model self.traced = set() - def _run(self, attn_metadata, input_ids, position_ids, logit_rows) -> torch.Tensor: + 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 [] - logits = self.model( - input_ids, position_ids, attn_metadata, logit_rows=logit_rows, trace=trace + 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 logits + return out class PrefillTarget(_Target): @@ -122,6 +181,46 @@ def forward( 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).""" @@ -139,11 +238,71 @@ def __init__(self, model_config: ModelConfig): 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" + 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 @@ -162,6 +321,83 @@ def load_checkpoint(self, checkpoint_dir: str) -> dict: 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, @@ -169,17 +405,47 @@ def forward( position_ids: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, return_context_logits: bool = False, + spec_metadata=None, + resource_manager=None, **kwargs, - ) -> torch.Tensor: + ): 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("spec_metadata") is not None or kwargs.get("lora_params") is not None: - raise ValueError("speculative decoding and LoRA are not supported by this target") + 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) - 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 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"] +__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 index fccde334b730..e5f9247cda6e 100644 --- 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 @@ -24,11 +24,15 @@ 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 [--port 8000] [--max-seq-len 69632] [--max-num-tokens 8192] \\ - [--max-batch-size 1] [--iter-stats] - -``--iter-stats`` enables the engine's iteration statistics (``GET /metrics``: per-step -scheduled, context and generation request counts, latency and memory). + --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 @@ -183,6 +187,25 @@ def load_tokenizer(checkpoint: str): 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, *, @@ -191,11 +214,14 @@ def llm_args( 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 @@ -204,7 +230,7 @@ def llm_args( 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}") - return { + args = { "model": checkpoint, "tokenizer": load_tokenizer(checkpoint), "backend": "pytorch", @@ -223,6 +249,9 @@ def llm_args( "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: @@ -239,6 +268,11 @@ def main(argv: Optional[list] = None) -> None: 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 @@ -256,6 +290,7 @@ def main(argv: Optional[list] = None) -> None: 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, 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 index de0bbfd9a96a..9199b09237f4 100644 --- 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 @@ -10,22 +10,31 @@ 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 -from typing import Any +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 from its checkpoint directory.""" + """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. @@ -35,16 +44,34 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, **kwargs) -> dict[ 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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/routing.py index e450443d603b..9c8a301a870d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/routing.py @@ -12,6 +12,11 @@ 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) @@ -20,6 +25,39 @@ def _quant_algo(pretrained_config: Any) -> Optional[str]: 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, @@ -32,10 +70,10 @@ def route( """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, no speculative - decoding. ``kv_cache_quant_algo`` is the effective engine setting, including runtime - overrides. No additional engine choice in ``model_config`` or ``llm_args`` selects a - different target. + 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, @@ -51,12 +89,13 @@ def route( 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 spec_config is None + and speculation is None and kv_cache_quant_algo is None ): from .full_model_nvfp4__sm_103__tp4.modeling import ModelingV2BeamFullModelNvfp4Sm103Tp4 @@ -66,8 +105,11 @@ def route( 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"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 or speculative decoding" + "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/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/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/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 91d3ff6d5605..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 @@ -7302,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/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 index 80a59f75a72e..5f5d2958e6f0 100644 --- 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 @@ -214,3 +214,146 @@ def test_load_pretrained_config_preserves_beam_rope(tmp_path): 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/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 index f8a621534b95..fddde53713ab 100644 --- 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 @@ -43,6 +43,7 @@ from __future__ import annotations import contextlib +import gc import math import types import zlib @@ -157,6 +158,15 @@ def token_stream(name: str, length: int) -> list: _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: 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 index 8b0eb5f3251c..dcb192a34dc3 100644 --- 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 @@ -32,6 +32,7 @@ from __future__ import annotations import contextlib +import gc import math import types import zlib @@ -205,6 +206,15 @@ def generate_ids(seed: int, length: int) -> torch.Tensor: _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: