diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 097edc9e2e6e..33e384dc3381 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -103,7 +103,7 @@ llm = LLM("/path/to/target_model", speculative_config=speculative_config, disabl ### MTP -MTP is currently only supported by Deepseek. MTP can be tuned with the following configuration options: +MTP is supported by DeepSeek models and other architectures that ship native MTP modules (including Step-3.x). MTP can be tuned with the following configuration options: * `max_draft_len`: Maximum draft candidate length. * `num_nextn_predict_layers`: Number of MTP modules to use. Currently must match `max_draft_len`. diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index e2b1f1d18602..802d324b5405 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -50,6 +50,7 @@ The following is a table of supported models for the PyTorch backend: | `SeedOssForCausalLM` [^5] | Seed OSS, Seed-Coder | `ByteDance-Seed/Seed-OSS-36B-Instruct` | | `SkyworkR1V2ForConditionalGeneration` [^5] | Skywork R1V2, Skywork SWE | `Skywork/Skywork-R1V2-38B` | | `SmolLM3ForCausalLM` [^5] | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | +| `Step3p7ForConditionalGeneration` [^8]| Step-3.7-Flash | `stepfun-ai/Step-3.7-Flash` | ## Model-Feature Support Matrix (Key Models) @@ -69,6 +70,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `Glm4MoeLiteForCausalLM` [^5] | Yes | Yes | Untested | Untested | Yes | No | No | No | No | Yes | Untested | Untested | N/A | Untested | Untested | | `NemotronHForCausalLM` (Super) | Yes | Yes | Untested | Untested | Yes | Yes | No | No | No | Yes | Yes | Untested | N/A | Untested | Untested | | `Gemma4ForConditionalGeneration` | Untested | Yes | Untested | No | Yes | No | No | No | No | Yes | Untested | No | Yes | Untested | Untested | +| `Step3p7ForConditionalGeneration`| Yes | Yes | Yes | Untested | Untested | Yes | No | No | No | Yes | Untested | Untested | Yes | Untested | Untested | [^1]: Chunked Prefill for MLA can only be enabled on SM100/SM103. [^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103 and in BF16/FP8 KV cache dtype. @@ -76,6 +78,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl [^5]: Supported via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend. See [AD Configs](../../../examples/auto_deploy/model_registry/configs). [^6]: Also supports text-only inference via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend. [^7]: Text-only support via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend. +[^8]: Supports text and image inputs. The vision tower runs in BF16 even when the text decoder is quantized (FP8 block-scale or NVFP4). The text decoder is also usable standalone (text-only) via the `Step3p5ForCausalLM` architecture. [^9]: Audio modality only supported on E2B/E4B variants. # Multimodal Feature Support Matrix (PyTorch Backend) @@ -96,6 +99,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `Qwen2_5_VLForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | L + I + V | | `Qwen3VLForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | L + I + V | | `Qwen3VLMoeForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | L + I + V | +| `Step3p7ForConditionalGeneration` | Yes | Yes | Untested | Yes | Untested | Untested | Untested | Untested | L + I | Note: - L: Language diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 9c3b032421b2..16d9875083b4 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -46,6 +46,8 @@ from .modeling_seedoss import SeedOssForCausalLM from .modeling_siglip import SiglipVisionModel from .modeling_starcoder2 import Starcoder2ForCausalLM +from .modeling_step3p7 import Step3p7ForCausalLM +from .modeling_step3p7vl import Step3p7VLForConditionalGeneration from .modeling_utils import get_model_architecture from .modeling_vila import VilaModel @@ -98,6 +100,8 @@ "Qwen3VLModel", "MiniMaxM2ForCausalLM", "Cohere2ForCausalLM", + "Step3p7ForCausalLM", + "Step3p7VLForConditionalGeneration", ] if transformers.__version__ >= "4.45.1": diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 62b24956d077..d960dd401eea 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1423,6 +1423,9 @@ def __init__( case "qwen3_next" | "qwen3_5_text" | "qwen3_5_moe_text": from .modeling_qwen3_next import Qwen3NextMTP mtp_layer = Qwen3NextMTP + case "step3p7" | "step3p5": + from .modeling_step3p7 import Step3p7MTP + mtp_layer = Step3p7MTP case _: raise ValueError( f"Model type {model_type} not supported for MTP") diff --git a/tensorrt_llm/_torch/models/modeling_step3p7.py b/tensorrt_llm/_torch/models/modeling_step3p7.py new file mode 100644 index 000000000000..ff3ba6bee2c6 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_step3p7.py @@ -0,0 +1,1689 @@ +# 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 +"""TensorRT-LLM PyTorch backend for the Step3p7 Flash text decoder. + +This module owns the text-only causal LM plus its MTP draft layers; the +sibling file ``modeling_step3p7vl`` (PerceptionEncoder vision tower + +multimodal input processor + ForConditionalGeneration wrapper) builds on +top of it for the VLM checkpoint variants. + +Step3p7 has three notable per-layer behaviors not shared with other models: + +- Per-layer RoPE: full-attention layers (0, 4, 8, ..., 44) use partial rotary + with llama3 scaling and theta 5e6; sliding-attention layers use full rotary + with theta 1e4 and no scaling. +- Gemma-style RMSNorm (``weight + 1``) on Q/K and all layer norms, and a + head-wise output gate (``sigmoid(g_proj(hidden))`` multiplied per attention + head before ``o_proj``). +- Layers 43/44 carry non-zero SwiGLU clamp limits for routed and shared + experts; routed experts use sigmoid+bias top-k routing with renormalization + and ``routed_scaling_factor=3.0``. + +Supported checkpoint variants: BF16 reference, FP8 block-scale flash, and +NVFP4 multimodal (with text decoder under ``model.language_model.*``, +flattened at load time). The on-disk checkpoints also carry 3 MTP layers +(loaded when ``MTPDecodingConfig`` is enabled) and a vision tower (ignored +here; handled by the VLM wrapper). +""" + +from __future__ import annotations + +import copy +import os +from typing import Dict, List, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + +from ..attention_backend import AttentionMetadata +from ..attention_backend.interface import ( + PositionalEmbeddingParams, + PredefinedAttentionMask, + RopeParams, +) +from ..distributed import AllReduce, AllReduceParams, allgather +from ..model_config import ModelConfig +from ..modules.attention import Attention +from ..modules.decoder_layer import DecoderLayer +from ..modules.embedding import Embedding, LMHead +from ..modules.fused_moe import create_moe +from ..modules.fused_moe.interface import MoEWeightLoadingMode +from ..modules.fused_moe.routing import MiniMaxM2MoeRoutingMethod +from ..modules.gated_mlp import GatedMLP +from ..modules.linear import Linear, TensorParallelMode +from ..modules.rms_norm import RMSNorm +from ..speculative import SpecMetadata +from ..utils import AuxStreamType, create_lm_head_tp_mapping +from .modeling_speculative import SpecDecOneEngineForCausalLM, _slice_spec_position_ids +from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model + +# --------------------------------------------------------------------------- +# FP8 / NVFP4 dequant helpers (used by the layer-43/44 SwiGLU-clamp path) +# --------------------------------------------------------------------------- + +_DEFAULT_FULL_ATTENTION_PERIOD = 4 +_DEFAULT_ROPE_THETA = 10000.0 +_FP8_BLOCK_SIZE = 128 +_PYTHON_FALLBACK_DISABLED_LAYER = 999 +_PYTHON_FALLBACK_THRESHOLD_ENV = "STEP3P7_PYTHON_FALLBACK_THRESHOLD" +_KV_SCALE_SUFFIXES = (".k_scale", ".v_scale") + + +def _fp8_block_dequant_3d( + weight_fp8: torch.Tensor, scale_inv: torch.Tensor, block: int = _FP8_BLOCK_SIZE +) -> torch.Tensor: + """Dequantise ``(E, M, K)`` FP8 e4m3 block-scale tensor to bf16.""" + _, M, K = weight_fp8.shape + scale = scale_inv.to(torch.float32).repeat_interleave(block, dim=-2) + scale = scale[..., :M, :].repeat_interleave(block, dim=-1)[..., :K] + return (weight_fp8.to(torch.float32) * scale).to(torch.bfloat16) + + +# e2m1 nibble lookup: the 16 representable values for the NVFP4 4-bit float +# (1 sign + 2 exponent + 1 mantissa). +_E2M1_VALUES = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, +) + + +def _ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +def _nvfp4_dequant_batched( + weight_uint8: torch.Tensor, + block_scale_fp8: torch.Tensor, + global_scale_fp32: torch.Tensor, + target_device: Optional[torch.device] = None, +) -> torch.Tensor: + """Dequantise an N-D NVFP4 weight tensor to bf16 in one batched op. + + ``weight_uint8`` is ``(*B, M, K/2)`` packed (two e2m1 nibbles per byte), + ``block_scale_fp8`` is ``(*B, M, K/16)`` (fp8_e4m3), and + ``global_scale_fp32`` is scalar or ``(B,)``. For checkpoint tensors that + live on CPU, pass ``target_device='cuda'`` -- the per-element ``lut[idx]`` + gather costs ~50ms per expert on CPU but <1ms on B200. + """ + device = target_device or weight_uint8.device + w = weight_uint8.to(device=device, non_blocking=True) + s1 = block_scale_fp8.to(device=device, non_blocking=True) + s2 = global_scale_fp32.to(device=device, non_blocking=True) + + K = w.shape[-1] * 2 + lut = _E2M1_VALUES.to(device=device) + + high = (w >> 4) & 0x0F + low = w & 0x0F + vals = torch.empty(*w.shape[:-1], K, dtype=torch.float32, device=device) + vals[..., 0::2] = lut[low.long()] + vals[..., 1::2] = lut[high.long()] + + s2_b = s2.to(torch.float32) + for _ in range(s1.dim() - s2_b.dim()): + s2_b = s2_b.unsqueeze(-1) + scale = (s1.to(torch.float32) * s2_b).unsqueeze(-1) + + vals = vals.view(*w.shape[:-1], K // 16, 16) * scale + return vals.view(*w.shape[:-1], K).to(torch.bfloat16) + + +def _nvfp4_stack_dequant( + weights, + base: str, + w_name: str, + expert_ids: list, + target_device: Optional[torch.device] = None, +) -> torch.Tensor: + """Stack and dequant NVFP4 experts ``expert_ids`` for ``{base}..``.""" + w_stack = torch.stack([weights[f"{base}.{e}.{w_name}.weight"] for e in expert_ids]) + s1_stack = torch.stack([weights[f"{base}.{e}.{w_name}.weight_scale"] for e in expert_ids]) + s2_stack = torch.stack( + [weights[f"{base}.{e}.{w_name}.weight_scale_2"].reshape([]) for e in expert_ids] + ) + return _nvfp4_dequant_batched(w_stack, s1_stack, s2_stack, target_device=target_device) + + +# --------------------------------------------------------------------------- +# Config normalization +# --------------------------------------------------------------------------- + + +def _normalize_torch_dtype(cfg) -> None: + """Normalize ``cfg.torch_dtype`` from a string (HF JSON form) to ``torch.dtype``.""" + dt = getattr(cfg, "torch_dtype", None) + if dt is None: + cfg.torch_dtype = torch.bfloat16 + return + if isinstance(dt, str): + mapped = getattr(torch, dt, None) + if isinstance(mapped, torch.dtype): + cfg.torch_dtype = mapped + return + try: + from transformers.utils import STR_TO_DTYPE # type: ignore + + cfg.torch_dtype = STR_TO_DTYPE.get(dt, torch.bfloat16) + except ImportError: + cfg.torch_dtype = torch.bfloat16 + + +def _mirror_step3p7_text_aliases(pretrained_config: PretrainedConfig) -> None: + """Promote Step3p7 ``text_config`` aliases the runtime expects on the top level. + + Complements ``ModelConfig.from_pretrained``'s generic ``_mirror_text_subconfig_attrs`` + by plugging in Step3p7-specific aliases (``num_key_value_heads`` from + ``num_attention_groups``, ``num_nextn_predict_layers``, etc.) and normalising + ``torch_dtype`` from the raw HF JSON string to ``torch.dtype``. + """ + text_config = getattr(pretrained_config, "text_config", None) + if text_config is None: + return + + _normalize_torch_dtype(pretrained_config) + _normalize_torch_dtype(text_config) + + if not hasattr(pretrained_config, "num_key_value_heads") and hasattr( + text_config, "num_attention_groups" + ): + pretrained_config.num_key_value_heads = text_config.num_attention_groups + if not hasattr(text_config, "num_key_value_heads") and hasattr( + text_config, "num_attention_groups" + ): + text_config.num_key_value_heads = text_config.num_attention_groups + + if not hasattr(pretrained_config, "max_position_embeddings") and hasattr( + text_config, "max_position_embeddings" + ): + pretrained_config.max_position_embeddings = text_config.max_position_embeddings + + if not hasattr(pretrained_config, "num_nextn_predict_layers") and hasattr( + text_config, "num_nextn_predict_layers" + ): + pretrained_config.num_nextn_predict_layers = text_config.num_nextn_predict_layers + + +def _get_text_config(model_config: ModelConfig) -> PretrainedConfig: + """Return the text sub-config, falling back to the top-level config.""" + cfg = model_config.pretrained_config + text_cfg = getattr(cfg, "text_config", None) + return text_cfg if text_cfg is not None else cfg + + +def _layer_attention_type(text_config: PretrainedConfig, layer_idx: int) -> str: + layer_types = getattr(text_config, "layer_types", None) or [] + if layer_idx < len(layer_types): + return layer_types[layer_idx] + # Default: layer 0,4,8,... are full attention. + if layer_idx % _DEFAULT_FULL_ATTENTION_PERIOD == 0: + return "full_attention" + return "sliding_attention" + + +def _layer_query_heads(text_config: PretrainedConfig, layer_idx: int) -> int: + if _layer_attention_type(text_config, layer_idx) == "sliding_attention": + other = getattr(text_config, "attention_other_setting", None) or {} + if "num_attention_heads" in other: + return int(other["num_attention_heads"]) + return int(text_config.num_attention_heads) + + +def _layer_kv_heads(text_config: PretrainedConfig, layer_idx: int) -> int: + if _layer_attention_type(text_config, layer_idx) == "sliding_attention": + other = getattr(text_config, "attention_other_setting", None) or {} + if "num_attention_groups" in other: + return int(other["num_attention_groups"]) + return int( + getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_groups")) + ) + + +def _per_layer_lookup(text_config: PretrainedConfig, values, layer_idx: int) -> float: + """Pick ``values[layer_idx]`` from a per-layer list, falling back to the most + recent entry that matches the current layer's attention type.""" + if layer_idx < len(values): + return float(values[layer_idx]) + layer_types = getattr(text_config, "layer_types", None) + if layer_types: + cur_type = _layer_attention_type(text_config, layer_idx) + for prev_idx in range(min(len(values), len(layer_types)) - 1, -1, -1): + if layer_types[prev_idx] == cur_type: + return float(values[prev_idx]) + return float(values[-1]) + + +def _layer_rope_theta(text_config: PretrainedConfig, layer_idx: int) -> float: + theta = getattr(text_config, "rope_theta", _DEFAULT_ROPE_THETA) + if isinstance(theta, (list, tuple)): + return _per_layer_lookup(text_config, theta, layer_idx) + return float(theta) + + +def _layer_partial_rotary(text_config: PretrainedConfig, layer_idx: int) -> float: + factors = getattr(text_config, "partial_rotary_factors", None) + if factors is None: + return 1.0 + return _per_layer_lookup(text_config, factors, layer_idx) + + +def _layer_uses_rope_scaling(text_config: PretrainedConfig, layer_idx: int) -> bool: + """llama3 scaling only applies to layer types listed in yarn_only_types.""" + yarn_only = getattr(text_config, "yarn_only_types", None) + if not yarn_only: + return True + return _layer_attention_type(text_config, layer_idx) in yarn_only + + +def _layer_swiglu_limit( + text_config: PretrainedConfig, layer_idx: int, shared: bool = False +) -> Optional[float]: + name = "swiglu_limits_shared" if shared else "swiglu_limits" + limits = getattr(text_config, name, None) + if limits is None or layer_idx >= len(limits): + return None + val = limits[layer_idx] + if val is None or float(val) == 0.0: + return None + return float(val) + + +def _is_moe_layer(text_config: PretrainedConfig, layer_idx: int) -> bool: + enum = getattr(text_config, "moe_layers_enum", None) + if enum is None: + return False + if isinstance(enum, str): + moe_layers = {int(x) for x in enum.split(",") if x.strip()} + else: + moe_layers = {int(x) for x in enum} + return layer_idx in moe_layers + + +def _parse_python_fallback_threshold() -> int: + try: + return int( + os.environ.get(_PYTHON_FALLBACK_THRESHOLD_ENV, str(_PYTHON_FALLBACK_DISABLED_LAYER)) + ) + except ValueError: + return _PYTHON_FALLBACK_DISABLED_LAYER + + +def _select_python_expert_path( + model_config: ModelConfig, + text_config: PretrainedConfig, + layer_idx: int, +) -> tuple[float | None, bool, str]: + """Return ``(swiglu_limit, use_python_path, reason)`` for a routed MoE layer. + + The Python expert loop is selected for clamp-active layers (no FP8 backend + on B200 supports ``swiglu_limit``), for BF16 checkpoints (the FP8 kernel + can't run against bf16 weights), and via an env-var diagnostic override. + """ + swiglu_limit = _layer_swiglu_limit(text_config, layer_idx, shared=False) + qc = getattr(model_config, "quant_config", None) + is_bf16_checkpoint = qc is None or getattr(qc, "quant_algo", None) is None + + if swiglu_limit is not None and swiglu_limit > 0: + return swiglu_limit, True, "clamp" + if layer_idx >= _parse_python_fallback_threshold(): + return swiglu_limit, True, "late_layer" + if is_bf16_checkpoint: + return swiglu_limit, True, "bf16_kernel_workaround" + return swiglu_limit, False, "" + + +# --------------------------------------------------------------------------- +# MoE routing +# --------------------------------------------------------------------------- + + +class Step3p7RouterBiasHolder(nn.Module): + """Standalone parameter holder for ``router_bias`` (mirrors MiniMaxM2). + + The holder is attached as ``moe.router_bias`` so the resulting parameter + path ``moe.router_bias.router_bias`` matches the HF source key + ``moe.router_bias`` exactly when the generic loader filters weights. + """ + + def __init__(self, num_experts: int): + super().__init__() + self.router_bias = nn.Parameter( + torch.empty((num_experts,), dtype=torch.float32), + requires_grad=False, + ) + + def load_weights(self, weights: List[Dict]): + assert len(weights) == 1 + w = weights[0] + if "" in w: + src = w[""] + elif "router_bias" in w: + src = w["router_bias"] + else: + (src,) = w.values() + self.router_bias.copy_(src[:].to(self.router_bias.dtype)) + + +class Step3p7MoeRoutingMethod(MiniMaxM2MoeRoutingMethod): + """Step3p7 routing: ``sigmoid -> add bias -> top-k -> renormalize -> scale``. + + Inherits from ``MiniMaxM2MoeRoutingMethod`` so the TRTLLMGen + ``_extract_routing_params`` helper recognises us via ``isinstance`` and + feeds the bias pointer to the kernel. The MiniMax2 C++ routing path + hard-codes ``routeScale = 1.0f`` (see ``runner.cu``), so + ``routed_scaling_factor`` is applied to the MoE output in + ``Step3p7MoE.forward`` instead of inside the kernel. + """ + + def __init__( + self, + top_k: int, + num_experts: int, + callable_router_bias, + routed_scaling_factor: float = 1.0, + output_dtype: torch.dtype = torch.float32, + ): + super().__init__( + top_k=top_k, + num_experts=num_experts, + callable_e_score_correction_bias=callable_router_bias, + output_dtype=output_dtype, + ) + self.routed_scaling_factor = float(routed_scaling_factor) + + def apply( + self, + router_logits: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + scores = torch.sigmoid(router_logits.to(torch.float32)) + scores_with_bias = scores + self.router_bias.unsqueeze(0) + _, topk_idx = torch.topk(scores_with_bias, k=self.top_k, dim=1) + topk_weights = torch.gather(scores, 1, topk_idx) + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) + if self.routed_scaling_factor != 1.0: + topk_weights = topk_weights.to(torch.float32) * self.routed_scaling_factor + topk_weights = topk_weights.to(self.output_dtype) + return topk_idx.to(torch.int32), topk_weights + + @property + def router_bias(self) -> torch.Tensor: + return self.callable_e_score_correction_bias() + + +# --------------------------------------------------------------------------- +# Attention +# --------------------------------------------------------------------------- + + +_LLAMA3_ROPE_PARAM_KEYS = frozenset( + { + "rope_type", + "factor", + "original_max_position_embeddings", + "low_freq_factor", + "high_freq_factor", + } +) + + +def _build_per_layer_config(text_config: PretrainedConfig, layer_idx: int) -> PretrainedConfig: + """Shallow-copy ``text_config`` with per-layer RoPE / head scalars applied. + + Step3p7 stores ``num_attention_heads``, ``rope_theta`` etc. as per-layer + lists, while TRT-LLM's ``RopeParams.from_config`` and ``Attention.__init__`` + read them as top-level scalars. Transformers 5.x also keeps the per-layer + list inside the ``rope_parameters`` dict; ``RopeParams.from_config`` does + ``config.update(rope_parameters)``, so we must materialise a per-layer + ``rope_parameters`` too — otherwise the scalar gets overwritten by the list. + """ + cfg = copy.copy(text_config) + cfg.num_attention_heads = _layer_query_heads(text_config, layer_idx) + cfg.num_key_value_heads = _layer_kv_heads(text_config, layer_idx) + theta_scalar = _layer_rope_theta(text_config, layer_idx) + cfg.rope_theta = theta_scalar + cfg.partial_rotary_factor = _layer_partial_rotary(text_config, layer_idx) + + src_rope_params = getattr(text_config, "rope_parameters", None) + if isinstance(src_rope_params, dict): + cfg.rope_parameters = {**src_rope_params, "rope_theta": theta_scalar} + + if not _layer_uses_rope_scaling(text_config, layer_idx): + cfg.rope_scaling = None + if isinstance(getattr(cfg, "rope_parameters", None), dict): + cfg.rope_parameters = { + k: v for k, v in cfg.rope_parameters.items() if k not in _LLAMA3_ROPE_PARAM_KEYS + } + return cfg + + +class Step3p7Attention(Attention): + """Per-layer Step3p7 attention. + + Owns its own Q/K/V/O shapes derived from the per-layer query-head count. + Q/K Gemma-style RMSNorm and head-wise output gate are applied module-side + so the production attention backend (FlashInfer + KVCacheManagerV2) only + sees standard QKV. + """ + + def __init__(self, model_config: ModelConfig, layer_idx: int): + text_config = _get_text_config(model_config) + per_layer_cfg = _build_per_layer_config(text_config, layer_idx) + self.text_config = text_config + self.layer_attention_type = _layer_attention_type(text_config, layer_idx) + self.sliding_window = ( + text_config.sliding_window if self.layer_attention_type == "sliding_attention" else None + ) + self.head_dim = int( + getattr( + text_config, "head_dim", text_config.hidden_size // text_config.num_attention_heads + ) + ) + self.use_head_wise_gate = bool(getattr(text_config, "use_head_wise_attn_gate", False)) + + rope_params = RopeParams.from_config(per_layer_cfg) + rope_params.theta = per_layer_cfg.rope_theta + if _layer_uses_rope_scaling(text_config, layer_idx): + rope_params.scale_type = RotaryScalingType.llama3 + else: + rope_params.scale_type = RotaryScalingType.none + rope_params.scale = 1.0 + rope_params.dim = int(self.head_dim * _layer_partial_rotary(text_config, layer_idx)) + + pos_embd_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=rope_params, + ) + + # Swap in the per-layer pretrained_config for the duration of the base + # constructor so it sees scalar (not per-layer-list) head/RoPE fields. + # ``ModelConfig`` is frozen but ``pretrained_config`` is exempt. + original_pretrained_config = model_config.pretrained_config + model_config.pretrained_config = per_layer_cfg + try: + super().__init__( + hidden_size=text_config.hidden_size, + num_attention_heads=per_layer_cfg.num_attention_heads, + num_key_value_heads=per_layer_cfg.num_key_value_heads, + max_position_embeddings=per_layer_cfg.max_position_embeddings, + bias=False, + pos_embd_params=pos_embd_params, + rope_fusion=True, + layer_idx=layer_idx, + dtype=text_config.torch_dtype, + dense_bias=False, + config=model_config, + head_dim=self.head_dim, + ) + finally: + model_config.pretrained_config = original_pretrained_config + + self.q_norm = RMSNorm( + hidden_size=self.head_dim, + eps=text_config.rms_norm_eps, + dtype=text_config.torch_dtype, + use_gemma=True, + ) + self.k_norm = RMSNorm( + hidden_size=self.head_dim, + eps=text_config.rms_norm_eps, + dtype=text_config.torch_dtype, + use_gemma=True, + ) + + if self.use_head_wise_gate: + self.g_proj = Linear( + in_features=text_config.hidden_size, + out_features=per_layer_cfg.num_attention_heads, + bias=False, + dtype=text_config.torch_dtype, + mapping=self.qkv_proj.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=False, + quant_config=None, + ) + + def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + q_view = q.reshape(-1, self.head_dim) + k_view = k.reshape(-1, self.head_dim) + q = self.q_norm(q_view).reshape(q.shape) + k = self.k_norm(k_view).reshape(k.shape) + return q, k + + def apply_rope( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + position_ids: torch.Tensor, + ): + # Split before QK-norm: Gemma RMSNorm needs the heads-separated layout. + q, k, v = self.split_qkv(q, k, v) + q, k = self.apply_qk_norm(q, k) + return super().apply_rope(q, k, v, position_ids) + + def forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + attention_mask=PredefinedAttentionMask.CAUSAL, + attention_window_size: Optional[int] = None, + **kwargs, + ) -> torch.Tensor: + """Step3p7 attention with module-side QK-norm + RoPE + head gate. + + Bypasses the base ``Attention.forward`` to apply the per-head output + gate (``sigmoid(g_proj(hidden))``) between the attention backend and + ``o_proj``. Helix CP, LoRA, and attention sinks are not plumbed here. + """ + effective_window = ( + attention_window_size if attention_window_size is not None else self.sliding_window + ) + + if ( + not self.rope_fusion + and getattr(attn_metadata, "is_spec_dec_dynamic_tree", False) + and getattr(attn_metadata, "use_spec_decoding", False) + and getattr(attn_metadata, "spec_decoding_position_offsets", None) is not None + and attn_metadata.spec_decoding_position_offsets.dim() == 1 + and position_ids is not None + ): + position_ids = self._adjust_position_ids_for_spec_dec( + position_ids.clone(), attn_metadata + ).clamp_min_(0) + + num_text_layers = int(getattr(self.text_config, "num_hidden_layers", 0)) + if position_ids is not None and self.layer_idx >= num_text_layers: + position_ids = position_ids.clamp_min(0) + + qkv = self.qkv_proj(hidden_states) + q, k, v = self.apply_rope(qkv, None, None, position_ids) + q, k, v = self.convert_qkv(q, k, v) + + attn_output = self.forward_impl( + q, + k, + v, + attn_metadata, + attention_mask, + effective_window, + kwargs.get("attention_mask_data"), + mrope_config=kwargs.get("mrope_config"), + attention_sinks=None, + has_lora=False, + ) + + if self.use_head_wise_gate: + # self.num_heads is already per-rank after TP sharding. + gate = self.g_proj(hidden_states) + orig_shape = attn_output.shape + attn_output = attn_output.view(*orig_shape[:-1], self.num_heads, self.head_dim) + attn_output = attn_output * gate.unsqueeze(-1).sigmoid() + attn_output = attn_output.view(*orig_shape) + + return self.o_proj(attn_output) + + +# --------------------------------------------------------------------------- +# MLP / MoE / Decoder layer +# --------------------------------------------------------------------------- + + +class ClampedGatedMLP(GatedMLP): + """Dense SwiGLU MLP with optional per-layer clamp on gate/up activations. + + When ``swiglu_limit`` is set, ``gate`` is clamped to ``[-inf, limit]`` and + ``up`` is clamped to ``[-limit, limit]`` before the elementwise multiply + (matching source ``Step3p7MLP``). Otherwise this is a standard ``GatedMLP``. + """ + + def __init__( + self, + model_config: ModelConfig, + layer_idx: int, + intermediate_size: int, + swiglu_limit: Optional[float], + is_shared_expert: bool = False, + ): + text_config = _get_text_config(model_config) + super().__init__( + hidden_size=text_config.hidden_size, + intermediate_size=intermediate_size, + bias=False, + dtype=text_config.torch_dtype, + config=model_config, + layer_idx=layer_idx, + # Shared expert defers all-reduce to the routed MoE call. + reduce_output=not is_shared_expert, + overridden_tp_size=None, + is_shared_expert=is_shared_expert, + ) + self.swiglu_limit = swiglu_limit + + def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + if self.swiglu_limit is None: + return super().forward(hidden_states, **kwargs) + gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1) + gate = torch.nn.functional.silu(gate).clamp(max=self.swiglu_limit) + up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) + return self.down_proj(gate * up) + + +class Step3p7MoE(nn.Module): + """Routed MoE block: router gate + bias-holder + routed experts. + + The shared expert is owned by the decoder layer (HF source stores it + under ``share_expert.*``, not ``moe.share_expert.*``). + """ + + _CLAMP_BUFFER_NAMES = ("_clamp_gate_proj", "_clamp_up_proj", "_clamp_down_proj") + + def __init__( + self, + model_config: ModelConfig, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ) -> None: + super().__init__() + text_config = _get_text_config(model_config) + self.layer_idx = layer_idx + self.hidden_size = text_config.hidden_size + self.num_experts = int(text_config.moe_num_experts) + self.top_k = int(text_config.moe_top_k) + self.moe_intermediate_size = int(text_config.moe_intermediate_size) + self.routed_scaling_factor = float(getattr(text_config, "moe_router_scaling_factor", 1.0)) + self.need_fp32_gate = bool(getattr(text_config, "need_fp32_gate", False)) + self.mapping = model_config.mapping + self.enable_attention_dp = self.mapping.enable_attention_dp + + gate_dtype = torch.float32 if self.need_fp32_gate else text_config.torch_dtype + self.gate = Linear( + in_features=self.hidden_size, + out_features=self.num_experts, + bias=False, + dtype=gate_dtype, + quant_config=None, + ) + + self.router_bias = Step3p7RouterBiasHolder(self.num_experts) + routing_method = Step3p7MoeRoutingMethod( + top_k=self.top_k, + num_experts=self.num_experts, + callable_router_bias=lambda: self.router_bias.router_bias, + routed_scaling_factor=self.routed_scaling_factor, + ) + + ( + self._routed_swiglu_limit, + self._use_python_clamp, + self._python_path_reason, + ) = _select_python_expert_path(model_config, text_config, layer_idx) + + self.experts = create_moe( + num_experts=self.num_experts, + routing_method=routing_method, + hidden_size=self.hidden_size, + intermediate_size=self.moe_intermediate_size, + aux_stream_dict=aux_stream_dict, + dtype=text_config.torch_dtype, + reduce_results=False, + model_config=model_config, + layer_idx=layer_idx, + weight_loading_mode=MoEWeightLoadingMode.VANILLA, + ) + + if self._use_python_clamp: + self._allocate_clamp_buffers(text_config.torch_dtype) + + self.allreduce = None + if not self.enable_attention_dp and self.mapping.tp_size > 1: + self.allreduce = AllReduce( + mapping=self.mapping, strategy=model_config.allreduce_strategy + ) + + def _allocate_clamp_buffers( + self, + dtype: torch.dtype, + shapes: Optional[Tuple[tuple, tuple, tuple]] = None, + device: Optional[torch.device] = None, + ) -> None: + """Allocate non-persistent expert buffers for the Python clamp path.""" + if shapes is None: + num_local = max(1, self.num_experts // max(1, self.mapping.moe_ep_size)) + gw_shape = (num_local, self.moe_intermediate_size, self.hidden_size) + up_shape = (num_local, self.moe_intermediate_size, self.hidden_size) + dn_shape = (num_local, self.hidden_size, self.moe_intermediate_size) + else: + gw_shape, up_shape, dn_shape = shapes + for name, shape in zip(self._CLAMP_BUFFER_NAMES, (gw_shape, up_shape, dn_shape)): + self.register_buffer( + name, torch.empty(shape, dtype=dtype, device=device), persistent=False + ) + self._clamp_num_local_experts = int(gw_shape[0]) + self._clamp_weights_loaded = False + + def _copy_clamp_weights(self, gate: torch.Tensor, up: torch.Tensor, down: torch.Tensor) -> None: + if self._clamp_gate_proj.shape != gate.shape: + self._allocate_clamp_buffers( + self._clamp_gate_proj.dtype, + shapes=(tuple(gate.shape), tuple(up.shape), tuple(down.shape)), + device=self._clamp_gate_proj.device, + ) + dev, dt = self._clamp_gate_proj.device, self._clamp_gate_proj.dtype + self._clamp_gate_proj.copy_(gate.to(device=dev, dtype=dt)) + self._clamp_up_proj.copy_(up.to(device=dev, dtype=dt)) + self._clamp_down_proj.copy_(down.to(device=dev, dtype=dt)) + self._clamp_weights_loaded = True + + def _clamp_weight_device(self) -> torch.device: + dev = self._clamp_gate_proj.device + if dev.type == "meta": + return torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + return dev + + def _local_expert_ids(self) -> list[int]: + ep_size = max(1, self.mapping.moe_ep_size) + num_local = max(1, self.num_experts // ep_size) + local_start = int(self.mapping.moe_ep_rank) * num_local + return list(range(local_start, local_start + num_local)) + + def load_clamp_weights_from_fp8_experts(self) -> None: + """Populate clamp buffers by dequantising the backend's expert weights. + + Called once by the model loader after ``super().load_weights`` has + materialised the backend tensors. For the bf16 reference checkpoint + the TRTLLMGen-Bf16 backend reshuffles ``w3_w1_weight`` into a 4D + BlockMajorK layout that can't be sliced back; in that case + ``Step3p7ForCausalLM._capture_bf16_clamp_weights`` runs earlier (from + the raw HF tensors) and this method short-circuits. + """ + if not self._use_python_clamp or self._clamp_weights_loaded: + return + # ``self.experts`` is a ConfigurableMoE wrapper; weights live on the backend. + e = getattr(self.experts, "backend", None) or self.experts + if not hasattr(e, "w3_w1_weight") or not hasattr(e, "w2_weight"): + return + w3_w1 = e.w3_w1_weight.data + w2 = e.w2_weight.data + moe_inter = self.moe_intermediate_size + # w3_w1 layout is [w3 (up), w1 (gate)] along dim 1, same for FP8 and bf16 backends. + w3 = w3_w1[:, :moe_inter, :] + w1 = w3_w1[:, moe_inter:, :] + scale_w3_w1 = getattr(e, "w3_w1_weight_scaling_factor", None) + scale_w2 = getattr(e, "w2_weight_scaling_factor", None) + is_fp8 = ( + scale_w3_w1 is not None and scale_w2 is not None and w3_w1.dtype == torch.float8_e4m3fn + ) + if is_fp8: + intermediate_blocks = _ceil_div(moe_inter, _FP8_BLOCK_SIZE) + w3_scale = scale_w3_w1.data[:, :intermediate_blocks, :] + w1_scale = scale_w3_w1.data[:, intermediate_blocks:, :] + up_bf16 = _fp8_block_dequant_3d(w3, w3_scale) + gate_bf16 = _fp8_block_dequant_3d(w1, w1_scale) + down_bf16 = _fp8_block_dequant_3d(w2, scale_w2.data) + else: + up_bf16 = w3.to(torch.bfloat16).contiguous() + gate_bf16 = w1.to(torch.bfloat16).contiguous() + down_bf16 = w2.to(torch.bfloat16).contiguous() + self._copy_clamp_weights(gate_bf16, up_bf16, down_bf16) + + def _python_clamped_moe_forward( + self, h: torch.Tensor, router_logits: torch.Tensor + ) -> torch.Tensor: + """Python expert path with explicit SwiGLU clamp. + + Mirrors HF source ``Step3p7MoEMLP.get_expert_output`` and produces a + per-rank partial sum (the decoder layer's all-reduce combines partials + across EP ranks). + + CUDA-graph capture safe: the per-expert loop has a fixed Python count, + and routing is materialised as a dense ``(N, num_experts)`` matrix via + ``scatter_add_`` so each iteration multiplies by a fixed-shape column. + Wasted compute is ``num_local / top_k`` over the masked version, + negligible at Step3p7 layer dimensions. + """ + topk_idx, topk_weights = self.experts.routing_method.apply(router_logits) + N = topk_idx.shape[0] + num_local = self._clamp_num_local_experts + local_start = self.mapping.moe_ep_rank * num_local + limit = float(self._routed_swiglu_limit or 0.0) + + weight_per_expert = torch.zeros((N, self.num_experts), dtype=torch.float32, device=h.device) + weight_per_expert.scatter_add_(1, topk_idx.long(), topk_weights.to(torch.float32)) + + # HF source computes experts in fp32 and casts each result back to + # bf16 before accumulation -- keep that rounding point. + x_f32 = h.to(torch.float32) + output = torch.zeros((N, self.hidden_size), dtype=h.dtype, device=h.device) + for local_e in range(num_local): + global_e = local_start + local_e + gate_w = self._clamp_gate_proj[local_e].to(torch.float32) + up_w = self._clamp_up_proj[local_e].to(torch.float32) + down_w = self._clamp_down_proj[local_e].to(torch.float32) + gate = torch.nn.functional.silu(x_f32 @ gate_w.t()) + up = x_f32 @ up_w.t() + if limit > 0.0: + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + expert_out = (gate * up) @ down_w.t() + expert_out = expert_out * weight_per_expert[:, global_e : global_e + 1] + output = output + expert_out.to(h.dtype) + return output + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams] = None, + **kwargs, + ) -> torch.Tensor: + assert hidden_states.shape[-1] == self.hidden_size + orig_shape = hidden_states.shape + h = hidden_states.view(-1, self.hidden_size) + + gate_input = h.to(torch.float32) if self.need_fp32_gate else h + router_logits = self.gate(gate_input) + + if self._use_python_clamp and self._clamp_weights_loaded: + # Python path's routing.apply already scales topk_weights. + return self._python_clamped_moe_forward(h, router_logits).view(orig_shape) + + routed = self.experts( + h, + router_logits, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens, + use_dp_padding=False, + ) + # TRTLLMGen MiniMax2 kernel hard-codes routeScale=1.0, so apply + # ``routed_scaling_factor`` to the MoE output here (mathematically + # equivalent to scaling each topk weight). + if self.routed_scaling_factor != 1.0: + routed = routed * self.routed_scaling_factor + return routed.view(orig_shape) + + +def _copy_model_config_with_quant(model_config: ModelConfig, quant_config: object) -> ModelConfig: + """Shallow-copy ``model_config`` with a replacement ``quant_config``. + + ``ModelConfig`` is frozen but ``quant_config`` is exempt; we also need to + clear ``quant_config_dict``, which requires bypassing the freeze. + """ + cloned = copy.copy(model_config) + cloned.quant_config = quant_config + object.__setattr__(cloned, "_frozen", False) + try: + cloned.quant_config_dict = None + finally: + object.__setattr__(cloned, "_frozen", True) + return cloned + + +def _model_config_without_quant(model_config: ModelConfig) -> ModelConfig: + """Shallow copy with a no-op quant config (for bf16 paths like attention).""" + from tensorrt_llm.models.modeling_utils import QuantConfig + + return _copy_model_config_with_quant(model_config, QuantConfig()) + + +class Step3p7DecoderLayer(DecoderLayer): + def __init__( + self, + model_config: ModelConfig, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + super().__init__() + text_config = _get_text_config(model_config) + self.layer_idx = layer_idx + self.hidden_size = text_config.hidden_size + + # Attention/dense MLP/shared expert stay bf16 (only routed experts + # are FP8/NVFP4). Attention still needs ``kv_cache_quant_algo`` so + # FP8 KV cache (set by modelopt) reaches the attention backend. + bf16_model_config = _model_config_without_quant(model_config) + attn_model_config = _model_config_keep_kv_quant(model_config) + self.self_attn = Step3p7Attention(attn_model_config, layer_idx) + + self.is_moe_layer = _is_moe_layer(text_config, layer_idx) + if self.is_moe_layer: + self.moe = Step3p7MoE( + model_config, layer_idx=layer_idx, aux_stream_dict=aux_stream_dict + ) + self.share_expert = ClampedGatedMLP( + bf16_model_config, + layer_idx=layer_idx, + intermediate_size=int( + getattr(text_config, "share_expert_dim", text_config.moe_intermediate_size) + ), + swiglu_limit=_layer_swiglu_limit(text_config, layer_idx, shared=True), + is_shared_expert=True, + ) + self.mlp = None + else: + self.mlp = ClampedGatedMLP( + bf16_model_config, + layer_idx=layer_idx, + intermediate_size=int(text_config.intermediate_size), + swiglu_limit=_layer_swiglu_limit(text_config, layer_idx, shared=False), + is_shared_expert=False, + ) + self.moe = None + self.share_expert = None + + # All Step3p7 RMSNorms (Q/K norms and every layer norm) use Gemma-style + # ``(weight + 1)`` scaling, not the standard ``weight`` form. + self.input_layernorm = RMSNorm( + hidden_size=text_config.hidden_size, + eps=text_config.rms_norm_eps, + dtype=text_config.torch_dtype, + use_gemma=True, + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=text_config.hidden_size, + eps=text_config.rms_norm_eps, + dtype=text_config.torch_dtype, + use_gemma=True, + ) + + self.allreduce = None + if ( + self.is_moe_layer + and not model_config.mapping.enable_attention_dp + and model_config.mapping.tp_size > 1 + ): + self.allreduce = AllReduce( + mapping=model_config.mapping, strategy=model_config.allreduce_strategy + ) + + def forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + residual: Optional[torch.Tensor], + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + **kwargs, + ) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + if self.moe is not None: + routed = self.moe(hidden_states, attn_metadata) + shared = self.share_expert(hidden_states) + hidden_states = routed + shared + if self.allreduce is not None: + hidden_states = self.allreduce(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +# --------------------------------------------------------------------------- +# Top-level model & registration +# --------------------------------------------------------------------------- + + +class Step3p7TextModel(DecoderModel): + def __init__(self, model_config: ModelConfig): + super().__init__(model_config) + text_config = _get_text_config(model_config) + self.vocab_size = int(text_config.vocab_size) + self.num_hidden_layers = int(text_config.num_hidden_layers) + self.aux_stream_dict = { + AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(), + AuxStreamType.MoeBalancer: torch.cuda.Stream(), + AuxStreamType.MoeOutputMemset: torch.cuda.Stream(), + } + + self.embed_tokens = Embedding( + self.vocab_size, + int(text_config.hidden_size), + dtype=text_config.torch_dtype, + ) + self.layers = nn.ModuleList( + [ + Step3p7DecoderLayer(model_config, idx, self.aux_stream_dict) + for idx in range(self.num_hidden_layers) + ] + ) + self.norm = RMSNorm( + hidden_size=int(text_config.hidden_size), + eps=text_config.rms_norm_eps, + dtype=text_config.torch_dtype, + use_gemma=True, + ) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds.") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + hidden_states = inputs_embeds + residual = None + for decoder_layer in self.layers[: self.num_hidden_layers]: + hidden_states, residual = decoder_layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=residual, + **kwargs, + ) + if residual is not None: + hidden_states = (hidden_states + residual).to(hidden_states.dtype) + return hidden_states + + +class Step3p7MTPHead(nn.Module): + """Step3p7 MTP shared head. + + The checkpoint stores a per-MTP-layer head under + ``model.layers..transformer.shared_head``. The MTP block returns the + residual-added hidden states, and this module applies ``norm`` before the + per-layer output projection to match the Step3p5/Step3p7 reference layout. + """ + + def __init__(self, model_config: ModelConfig): + super().__init__() + text_config = _get_text_config(model_config) + self.model_config = model_config + self.norm = RMSNorm( + hidden_size=int(text_config.hidden_size), + eps=text_config.rms_norm_eps, + dtype=text_config.torch_dtype, + use_gemma=True, + ) + mapping = model_config.mapping + if mapping.enable_attention_dp and not mapping.enable_lm_head_tp_in_adp: + self.output = LMHead( + int(text_config.vocab_size), + int(text_config.hidden_size), + dtype=text_config.torch_dtype, + ) + else: + self.output = LMHead( + int(text_config.vocab_size), + int(text_config.hidden_size), + dtype=text_config.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + reduce_output=False, + ) + self.mapping_lm_head_tp = None + + def _get_last_token_states( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + last_tokens = torch.cumsum(attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 + return hidden_states[last_tokens] + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: LMHead, + attn_metadata: AttentionMetadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + del lm_head + if not return_context_logits: + if attn_metadata is not None: + hidden_states = self._get_last_token_states(hidden_states, attn_metadata) + else: + hidden_states = hidden_states[-1].unsqueeze(0) + + mapping = self.model_config.mapping + enable_attention_dp = mapping.enable_attention_dp + enable_lm_head_tp_in_adp = enable_attention_dp and mapping.enable_lm_head_tp_in_adp + + if enable_lm_head_tp_in_adp: + self.mapping_lm_head_tp = create_lm_head_tp_mapping(mapping, hidden_states.shape[0]) + hidden_states = allgather(hidden_states, self.mapping_lm_head_tp, dim=0) + + hidden_states = self.norm(hidden_states) + + if not enable_attention_dp or enable_lm_head_tp_in_adp: + self.output.gather_output = False + logits = self.output( + hidden_states, + mapping_lm_head_tp=self.mapping_lm_head_tp, + is_spec_decoding_head=True, + ) + if not enable_attention_dp or enable_lm_head_tp_in_adp: + self.output.gather_output = True + return logits + + +class Step3p7MTP(nn.Module): + """Step3p7 native MTP predictor layer. + + This mirrors vLLM's Step3p5-MTP layout for this checkpoint: + ``enorm(embed)``, ``hnorm(target_hidden)``, ``eh_proj(concat(...))``, a + regular Step3p7 decoder block under ``mtp_block``, then + a residual add. ``shared_head.norm`` is applied by ``Step3p7MTPHead`` when + draft logits are requested. + """ + + def __init__( + self, + model_config: ModelConfig, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + super().__init__() + text_config = _get_text_config(model_config) + self.model_config = model_config + self.enorm = RMSNorm( + hidden_size=int(text_config.hidden_size), + eps=text_config.rms_norm_eps, + dtype=text_config.torch_dtype, + use_gemma=True, + ) + self.hnorm = RMSNorm( + hidden_size=int(text_config.hidden_size), + eps=text_config.rms_norm_eps, + dtype=text_config.torch_dtype, + use_gemma=True, + ) + if model_config.mapping.enable_attention_dp: + self.eh_proj = Linear( + in_features=int(text_config.hidden_size) * 2, + out_features=int(text_config.hidden_size), + bias=False, + dtype=text_config.torch_dtype, + quant_config=None, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + else: + self.eh_proj = Linear( + in_features=int(text_config.hidden_size) * 2, + out_features=int(text_config.hidden_size), + bias=False, + dtype=text_config.torch_dtype, + tensor_parallel_mode=TensorParallelMode.ROW, + mapping=model_config.mapping, + reduce_output=True, + quant_config=None, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + self.mtp_block = Step3p7DecoderLayer(model_config, layer_idx, aux_stream_dict) + self.shared_head = Step3p7MTPHead(model_config) + + def forward( + self, + input_ids: torch.IntTensor, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + embed_tokens: Embedding, + attn_metadata: AttentionMetadata, + all_rank_num_tokens: Optional[List[int]] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + del all_rank_num_tokens + inputs_embeds = self.enorm(embed_tokens(input_ids)) + hidden_states = self.hnorm(hidden_states) + hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1) + + mapping = self.model_config.mapping + if mapping.tp_size > 1 and not mapping.enable_attention_dp: + hidden_states = torch.chunk(hidden_states, mapping.tp_size, dim=-1)[mapping.tp_rank] + hidden_states = self.eh_proj(hidden_states) + + hidden_states, residual = self.mtp_block( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=None, + spec_metadata=spec_metadata, + **kwargs, + ) + if residual is not None: + hidden_states = (hidden_states + residual).to(hidden_states.dtype) + if spec_metadata is not None: + spec_metadata.maybe_capture_hidden_states(0, hidden_states, None) + return hidden_states + + +_STACKED_MOE_PROJ_TO_W = { + "gate_proj": "w1", + "up_proj": "w3", + "down_proj": "w2", +} + +# Per-expert suffixes that may accompany the stacked weights — covers FP8 +# block-scale (``weight_scale_inv``) and NVFP4 (``weight_scale``, +# ``weight_scale_2``, ``input_scale``) checkpoints. +_STACKED_MOE_WEIGHT_SUFFIXES = ( + "weight", + "weight_scale_inv", + "weight_scale", + "weight_scale_2", + "input_scale", +) + + +def split_stacked_moe_weights(weights, text_config: PretrainedConfig) -> int: + """Expand stacked routed-expert tensors into per-expert keys, in place. + + HF source stores each projection as a single ``(num_experts, ...)`` tensor; + the VANILLA MoE loader expects per-expert keys + ``experts..w{1,2,3}.``. Slices along dim 0 (zero-copy view) and + deletes the stacked keys. Returns the number of MoE layers split. + """ + num_layers = int(text_config.num_hidden_layers) + num_experts = int(text_config.moe_num_experts) + + layers_split = 0 + for layer_idx in range(num_layers): + if not _is_moe_layer(text_config, layer_idx): + continue + prefix = f"model.layers.{layer_idx}.moe." + proj_present = [p for p in _STACKED_MOE_PROJ_TO_W if f"{prefix}{p}.weight" in weights] + if not proj_present: + continue + layers_split += 1 + + for proj in proj_present: + dst_w = _STACKED_MOE_PROJ_TO_W[proj] + for suffix in _STACKED_MOE_WEIGHT_SUFFIXES: + stacked_key = f"{prefix}{proj}.{suffix}" + if stacked_key not in weights: + continue + stacked = weights[stacked_key] + if stacked.shape[0] != num_experts: + raise RuntimeError( + f"Step3p7 stacked MoE tensor {stacked_key} has " + f"leading dim {stacked.shape[0]} but num_experts={num_experts}." + ) + for expert_id in range(num_experts): + weights[f"{prefix}experts.{expert_id}.{dst_w}.{suffix}"] = stacked[expert_id] + del weights[stacked_key] + return layers_split + + +def _prepare_step3p7_mtp_spec_config(model_config: ModelConfig) -> None: + """Populate MTP config fields before ``SpecDecOneEngineForCausalLM`` builds draft layers.""" + spec_config = getattr(model_config, "spec_config", None) + if spec_config is None or getattr(spec_config, "decoding_type", None) != "MTP": + return + model_layers = int(getattr(model_config.pretrained_config, "num_nextn_predict_layers", 0) or 0) + if model_layers <= 0: + return + + spec_config.num_nextn_predict_layers = model_layers + if not spec_config.spec_dec_mode.is_mtp_vanilla(): + return + + # max_draft_len is None when the user didn't set it (use the model default); + # otherwise cap it at the model's MTP layer count. + if spec_config.max_draft_len is None or spec_config.max_draft_len >= model_layers: + spec_config.max_draft_len = model_layers + spec_config.max_total_draft_tokens = spec_config.max_draft_len + + +_MTP_DIRECT_WEIGHT_PREFIXES = ( + "enorm.", + "hnorm.", + "eh_proj.", + "shared_head.", +) + + +# Multimodal Step3p7 checkpoints (e.g. NVFP4 NIM/VNIM) nest the text decoder +# under ``model.language_model.*`` instead of ``model.*``; flatten at load time. +_LANGUAGE_MODEL_RENAMES = ( + ("model.language_model.", "model."), + ("model.vision_model.", "vision_model."), + ("model.vit_large_projector", "vit_large_projector"), +) + + +def rewrite_language_model_keys(weights) -> int: + """Flatten the multimodal ``model.language_model.*`` namespace in place. + + Returns the number of keys rewritten. Idempotent and safe on text-only + checkpoints that already use ``model.*`` keys. + """ + if weights is None or not hasattr(weights, "keys"): + return 0 + rewritten = 0 + for key in list(weights.keys()): + for src_prefix, dst_prefix in _LANGUAGE_MODEL_RENAMES: + if key.startswith(src_prefix): + new_key = dst_prefix + key[len(src_prefix) :] + if new_key == key: + break + value = weights[key] + del weights[key] + weights[new_key] = value + rewritten += 1 + break + return rewritten + + +def strip_language_model_prefix_from_exclude_modules(exclude_modules): + """Drop ``language_model.`` from quant-config exclude patterns. + + ModelOpt's ``hf_quant_config.json`` for the multimodal NVFP4 checkpoint + expresses ``exclude_modules`` against the on-disk namespace; TRT-LLM matches + against the runtime module path, so we strip the prefix. ``re:`` entries are + passed through untouched (they opt out of glob processing). + """ + if exclude_modules is None: + return None + return [ + e.replace("model.language_model.", "model.") + if isinstance(e, str) and not e.startswith("re:") + else e + for e in exclude_modules + ] + + +def _model_config_keep_kv_quant(model_config: ModelConfig) -> ModelConfig: + """Strip weight quant but preserve ``kv_cache_quant_algo`` for attention.""" + from tensorrt_llm.models.modeling_utils import QuantConfig + + src_kv = getattr(getattr(model_config, "quant_config", None), "kv_cache_quant_algo", None) + return _copy_model_config_with_quant(model_config, QuantConfig(kv_cache_quant_algo=src_kv)) + + +def rewrite_mtp_weights_for_step3p7(weights, text_config: PretrainedConfig) -> int: + """Rewrite Step3p7 MTP checkpoint keys to the TRT-LLM module layout. + + Source stores the transformer block directly under ``model.layers.`` + and the head under ``transformer.shared_head``; TRT-LLM uses ``mtp_block`` + and ``shared_head`` (matching the Step3p5-MTP layout used by vLLM). + """ + if weights is None or not hasattr(weights, "keys"): + return 0 + + num_layers = int(text_config.num_hidden_layers) + num_mtp_layers = int(getattr(text_config, "num_nextn_predict_layers", 0) or 0) + if num_mtp_layers <= 0: + return 0 + + rewritten = 0 + keys_snapshot = list(weights.keys()) + for layer_idx in range(num_layers, num_layers + num_mtp_layers): + prefix = f"model.layers.{layer_idx}." + for key in keys_snapshot: + if not key.startswith(prefix): + continue + suffix = key[len(prefix) :] + if suffix.startswith("transformer.shared_head."): + new_key = prefix + suffix.replace("transformer.", "", 1) + elif suffix.startswith(_MTP_DIRECT_WEIGHT_PREFIXES): + continue + else: + new_key = prefix + "mtp_block." + suffix + if new_key == key: + continue + value = weights[key] + del weights[key] + weights[new_key] = value + rewritten += 1 + return rewritten + + +@register_auto_model("Step3p5ForCausalLM") +class Step3p7ForCausalLM(SpecDecOneEngineForCausalLM[Step3p7TextModel, PretrainedConfig]): + """Step3p7 text-only causal LM core. + + Registered under ``Step3p5ForCausalLM`` (the text-config architecture); + the VLM entry point ``Step3p7ForConditionalGeneration`` wraps this class + in ``Step3p7VLForConditionalGeneration`` (see ``modeling_step3p7vl``). + When the wrapper sees a request without multimodal data, it delegates + straight here so plain text generation (GSM8K-style) keeps the original + behaviour. MTP layers (45..47) are still loaded only when + ``MTPDecodingConfig`` enables one-engine MTP; the vision tower lives on + the wrapper. + """ + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids=input_ids, + attn_metadata=attn_metadata, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **kwargs, + ) + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + + normed_hidden_states = self.model.norm(hidden_states) + + if self.spec_worker is not None: + logits = self.logits_processor.forward( + normed_hidden_states[spec_metadata.gather_ids], + self.lm_head, + attn_metadata, + True, + ) + + spec_input_ids = input_ids + spec_position_ids = position_ids + if attn_metadata.padded_num_tokens is not None: + if input_ids is not None: + spec_input_ids = input_ids[: attn_metadata.num_tokens] + if position_ids is not None: + spec_position_ids = _slice_spec_position_ids( + position_ids, attn_metadata.num_tokens + ) + + return self.spec_worker( + input_ids=spec_input_ids, + position_ids=spec_position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) + + return self.logits_processor.forward( + normed_hidden_states, + self.lm_head, + attn_metadata, + return_context_logits, + ) + + def __init__(self, model_config: ModelConfig): + _mirror_step3p7_text_aliases(model_config.pretrained_config) + + qc = getattr(model_config, "quant_config", None) + if qc is not None and qc.exclude_modules is not None: + qc.exclude_modules = strip_language_model_prefix_from_exclude_modules( + qc.exclude_modules + ) + + # ``model_loader.py`` snapshots ``pretrained_config.torch_dtype`` into + # ``extra_attrs['allreduce_dtype']`` before construction; when HF 5.x + # keeps it as a string (trust_remote_code path), refresh the snapshot + # so the NCCL window prealloc path sees a real ``torch.dtype``. + if hasattr(model_config, "extra_attrs") and isinstance(model_config.extra_attrs, dict): + cur = model_config.extra_attrs.get("allreduce_dtype") + if isinstance(cur, str): + mapped = getattr(torch, cur, None) + model_config.extra_attrs["allreduce_dtype"] = ( + mapped if isinstance(mapped, torch.dtype) else torch.bfloat16 + ) + + _prepare_step3p7_mtp_spec_config(model_config) + text_config = _get_text_config(model_config) + super().__init__(Step3p7TextModel(model_config), model_config) + self.text_config = text_config + self.num_hidden_layers = int(text_config.num_hidden_layers) + + active_mtp_layers = 0 + if ( + model_config.spec_config is not None + and model_config.spec_config.spec_dec_mode.is_mtp_one_model() + ): + self.model.layers.extend(self.draft_model.mtp_layers) + self.epilogue.extend(self.draft_model.mtp_layers) + self.epilogue.append(self.spec_worker) + active_mtp_layers = len(self.draft_model.mtp_layers) + + num_mtp = int(getattr(text_config, "num_nextn_predict_layers", 0)) + mtp_prefixes = tuple( + f"model.layers.{idx}." + for idx in range(self.num_hidden_layers, self.num_hidden_layers + num_mtp) + ) + self.ignored_key_prefixes = ( + "vision_model.", + "vit_large_projector.", + *mtp_prefixes[active_mtp_layers:], + ) + + def get_ignored_key_prefixes(self) -> Tuple[str, ...]: + return self.ignored_key_prefixes + + def _capture_bf16_clamp_weights(self, weights) -> List[int]: + """Pre-load 2D expert weights for Python-clamp layers, before backend layout transforms. + + The bf16 reference checkpoint's TRTLLMGen-Bf16 backend reshuffles + ``w3_w1_weight`` into a 4D BlockMajorK layout during + ``post_load_weights``; the NVFP4 backend likewise transforms its + weights. The Python clamp path needs the original 2D ``(I, H)`` + per-expert tensors, so we capture them here before + ``super().load_weights`` runs. FP8 layers are skipped because the FP8 + backend keeps a 3D layout that ``load_clamp_weights_from_fp8_experts`` + can dequantize directly post-load. + + Returns the list of layer indices captured (for diagnostic logging). + """ + if weights is None or not hasattr(weights, "__getitem__"): + return [] + captured: List[int] = [] + for layer in self.model.layers: + moe = getattr(layer, "moe", None) + if moe is None or not getattr(moe, "_use_python_clamp", False): + continue + if getattr(moe, "_clamp_weights_loaded", False): + continue + probe_key = f"model.layers.{moe.layer_idx}.moe.experts.0.w1.weight" + if probe_key not in weights: + continue + probe = weights[probe_key] + if probe.dtype == torch.float8_e4m3fn: + continue + try: + local_ids = moe._local_expert_ids() + base = f"model.layers.{moe.layer_idx}.moe.experts" + if probe.dtype == torch.uint8: + # NVFP4: batch the per-element ``lut[idx]`` gather on GPU + # (orders of magnitude faster than CPU). + target_dev = moe._clamp_weight_device() + gate_stack = _nvfp4_stack_dequant(weights, base, "w1", local_ids, target_dev) + up_stack = _nvfp4_stack_dequant(weights, base, "w3", local_ids, target_dev) + down_stack = _nvfp4_stack_dequant(weights, base, "w2", local_ids, target_dev) + else: + gate_stack = torch.stack([weights[f"{base}.{e}.w1.weight"] for e in local_ids]) + up_stack = torch.stack([weights[f"{base}.{e}.w3.weight"] for e in local_ids]) + down_stack = torch.stack([weights[f"{base}.{e}.w2.weight"] for e in local_ids]) + except KeyError: + continue + moe._copy_clamp_weights(gate_stack, up_stack, down_stack) + captured.append(moe.layer_idx) + return captured + + def load_weights( + self, + weights, + weight_mapper=None, + skip_modules=None, + params_map=None, + allow_partial_loading: bool = False, + ): + """Step3p7 weight loader. + + Pre-processing before delegating to the generic loader: + 1. Flatten multimodal ``model.language_model.*`` keys to ``model.*``. + 2. Drop ignored prefixes (vision tower, inactive MTP) and FP8 KV scales. + 3. Rewrite active MTP keys to the TRT-LLM ``mtp_block`` / ``shared_head`` layout. + 4. Split stacked routed-expert tensors into per-expert keys. + 5. Capture 2D bf16/NVFP4 expert weights for Python-clamp layers. + """ + from tensorrt_llm.logger import logger as _logger + + skip_modules = list(skip_modules) if skip_modules else [] + rewrite_language_model_keys(weights) + + if hasattr(weights, "keys"): + for k in list(weights.keys()): + drop = any(k.startswith(p) for p in self.ignored_key_prefixes) or k.endswith( + _KV_SCALE_SUFFIXES + ) + if drop: + try: + del weights[k] + except KeyError: + pass + + rewrite_mtp_weights_for_step3p7(weights, self.text_config) + split_stacked_moe_weights(weights, self.text_config) + self._capture_bf16_clamp_weights(weights) + + rc = DecoderModelForCausalLM.load_weights( + self, + weights, + weight_mapper=weight_mapper, + skip_modules=[*skip_modules, "draft_model"], + params_map=params_map, + allow_partial_loading=allow_partial_loading, + ) + # Must run after the backend materialises w3_w1/w2. + for layer in self.model.layers: + moe = getattr(layer, "moe", None) + if moe is None or not getattr(moe, "_use_python_clamp", False): + continue + try: + moe.load_clamp_weights_from_fp8_experts() + except (AttributeError, KeyError, RuntimeError, ValueError) as e: + _logger.warning( + "[Step3p7] failed to populate bf16 expert weights for layer %d (%s): %s. " + "Forward will fall back to the FP8 backend without the Python path.", + moe.layer_idx, + getattr(moe, "_python_path_reason", ""), + str(e)[:256], + ) + return rc diff --git a/tensorrt_llm/_torch/models/modeling_step3p7vl.py b/tensorrt_llm/_torch/models/modeling_step3p7vl.py new file mode 100644 index 000000000000..7f4027b309f4 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_step3p7vl.py @@ -0,0 +1,1020 @@ +# 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 +"""Multimodal entry point for the Step3p7 Flash checkpoint. + +Wraps the text-only ``Step3p7ForCausalLM`` (text decoder + MTP) plus a +Perception-Encoder vision tower. When a request has no multimodal payload +this wrapper is a thin passthrough so plain text generation keeps the +original Step3p7 text path. + +The vision tower mirrors the HF reference checkpoint's ``vision_encoder.py``: +``StepRoboticsVisionEncoder`` is a Perception-Encoder-style ViT (patch +embedding via Conv2d, 47 transformer blocks with 2D RoPE, two trailing +Conv2d downsamplers). The matching projector ``vit_large_projector`` is a +single bf16 Linear from ``4 * width`` to ``text_config.hidden_size``. + +The vision tower is intentionally kept in raw torch (SDPA for non-causal +attention) instead of TensorRT-LLM's ``Attention`` module. ``Attention`` +specialises for causal/MLA text decoders; a faithful port of the HF code +keeps weight names trivially compatible (``vision_model.transformer.resblocks. +.attn.{in_proj_weight,in_proj_bias,out_proj.{weight,bias}}``) and avoids +plumbing the vision attention through the text attention metadata. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from PIL import Image +from torch import nn +from transformers import AutoProcessor, AutoTokenizer, PretrainedConfig +from transformers.activations import ACT2FN +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from tensorrt_llm.inputs.multimodal import MultimodalParams + +from ..._utils import nvtx_range +from ...inputs import ( + BaseMultimodalInputProcessor, + ContentFormat, + ExtraProcessedInputs, + MultimodalPlaceholderMetadata, + TextPrompt, + register_input_processor, +) +from ...logger import logger +from ...sampling_params import SamplingParams +from ..attention_backend import AttentionMetadata +from ..model_config import ModelConfig +from ..modules.layer_norm import LayerNorm +from ..speculative import SpecMetadata +from .modeling_multimodal_utils import ( + _is_disagg, + find_input_mm_embeds, + fuse_input_embeds, + get_multimodal_embeddings, +) +from .modeling_step3p7 import ( + Step3p7ForCausalLM, + _get_text_config, + _mirror_step3p7_text_aliases, + _normalize_torch_dtype, +) +from .modeling_utils import register_auto_model, register_vision_encoder + +# --------------------------------------------------------------------------- +# Vision tower +# --------------------------------------------------------------------------- + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotate last dim halves -- 2D RoPE helper used by the vision tower.""" + x = x.reshape(*x.shape[:-1], -1, 2) + x1, x2 = x.unbind(dim=-1) + x = torch.stack((-x2, x1), dim=-1) + return x.reshape(*x.shape[:-2], -1) + + +def _apply_rotary_emb(freqs: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + dtype = t.dtype + rot_dim = freqs.shape[-1] + t_rot = t[..., :rot_dim] + t_pass = t[..., rot_dim:] + t_rot = (t_rot * freqs.cos()) + (_rotate_half(t_rot) * freqs.sin()) + return torch.cat((t_rot, t_pass), dim=-1).to(dtype) + + +class Step3VisionRope2D(nn.Module): + """Cached 2D rotary positional embedding for the vision tower.""" + + def __init__( + self, + dim: int, + max_grid_height: int, + max_grid_width: int, + use_cls_token: bool = False, + theta: float = 10000.0, + theta_rescale_factor: float = 1.0, + ): + super().__init__() + self.dim = dim + self.max_grid_height = max_grid_height + self.max_grid_width = max_grid_width + self.use_cls_token = use_cls_token + self.theta = theta * theta_rescale_factor ** (dim / (dim - 2)) + self.register_buffer("freqs_cache", self._compute_2d_freqs(), persistent=False) + + def _compute_inv_freq(self, base: float, dim: int) -> torch.Tensor: + return 1.0 / (base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + + def _compute_freqs(self, t: torch.Tensor, inv_freq: torch.Tensor) -> torch.Tensor: + freqs = torch.einsum("..., f -> ... f", t.type(inv_freq.dtype), inv_freq) + return freqs.repeat_interleave(2, dim=-1) + + def _compute_2d_freqs(self) -> torch.Tensor: + grid_h = torch.arange(self.max_grid_height, dtype=torch.float) + grid_w = torch.arange(self.max_grid_width, dtype=torch.float) + if self.use_cls_token: + grid_h = grid_h + 1 + grid_w = grid_w + 1 + inv_freq = self._compute_inv_freq(self.theta, self.dim // 2) + freqs_h = self._compute_freqs(grid_h, inv_freq)[:, None].expand( + self.max_grid_height, self.max_grid_width, -1 + ) + freqs_w = self._compute_freqs(grid_w, inv_freq)[None, :].expand( + self.max_grid_height, self.max_grid_width, -1 + ) + freqs = torch.cat([freqs_w, freqs_h], dim=-1).reshape( + self.max_grid_height * self.max_grid_width, -1 + ) + if self.use_cls_token: + freqs = torch.cat([torch.zeros(1, freqs.shape[-1]), freqs], dim=0) + return freqs[None, None, ...] + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + grid_hw: Tuple[int, int], + ) -> Tuple[torch.Tensor, torch.Tensor]: + if grid_hw[0] != self.max_grid_height or grid_hw[1] != self.max_grid_width: + rows = torch.arange(grid_hw[0], device=q.device).view(-1, 1) + cols = torch.arange(grid_hw[1], device=q.device).view(1, -1) + positions = (rows * self.max_grid_width + cols).reshape(-1).to(torch.long) + if self.use_cls_token: + positions = torch.cat( + [ + torch.zeros(1, device=q.device, dtype=torch.long), + positions + 1, + ], + dim=0, + ) + freqs = self.freqs_cache.index_select(2, positions) + else: + freqs = self.freqs_cache + return _apply_rotary_emb(freqs, q), _apply_rotary_emb(freqs, k) + + +class Step3VisionLayerScale(nn.Module): + """Per-channel residual scaling used when ``ls_init_value`` is set.""" + + def __init__(self, dim: int, init_value: float): + super().__init__() + self.gamma = nn.Parameter(torch.full((dim,), init_value)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states * self.gamma + + +class Step3VisionMLP(nn.Module): + """``c_fc -> act -> c_proj`` FFN matching the HF weight names.""" + + def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str): + super().__init__() + self.c_fc = nn.Linear(hidden_size, intermediate_size, bias=True) + self.act_fn = ACT2FN[hidden_act] + self.c_proj = nn.Linear(intermediate_size, hidden_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.c_proj(self.act_fn(self.c_fc(hidden_states))) + + +class Step3VisionAttention(nn.Module): + """Vision MHA with 2D RoPE. + + HF stores the fused QKV projection as ``in_proj_weight``/``in_proj_bias`` + (a single ``3*H x H`` matrix); we keep the same parameter names so the + checkpoint loads without remapping. + """ + + def __init__( + self, + hidden_size: int, + num_heads: int, + max_grid_height: int, + max_grid_width: int, + use_cls_token: bool, + use_rope2d: bool, + rope_theta: float = 10000.0, + rope_theta_rescale_factor: float = 1.0, + ): + super().__init__() + if hidden_size % num_heads != 0: + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by num_heads ({num_heads})." + ) + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.scale = self.head_dim**-0.5 + + # HF parameter names: in_proj_weight (3H, H), in_proj_bias (3H,), out_proj. + self.in_proj_weight = nn.Parameter(torch.zeros(hidden_size * 3, hidden_size)) + self.in_proj_bias = nn.Parameter(torch.zeros(hidden_size * 3)) + self.out_proj = nn.Linear(hidden_size, hidden_size, bias=True) + + self.rope: Optional[Step3VisionRope2D] = None + if use_rope2d: + self.rope = Step3VisionRope2D( + dim=self.head_dim, + max_grid_height=max_grid_height, + max_grid_width=max_grid_width, + use_cls_token=use_cls_token, + theta=rope_theta, + theta_rescale_factor=rope_theta_rescale_factor, + ) + + def forward( + self, + hidden_states: torch.Tensor, + grid_hw: Tuple[int, int], + ) -> torch.Tensor: + # TODO: port the vision attention/projector to TRT-LLM modules + bsz, seq_len, _ = hidden_states.shape + qkv = F.linear(hidden_states, self.in_proj_weight, self.in_proj_bias) + q, k, v = qkv.chunk(3, dim=-1) + q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + v = v.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + if self.rope is not None: + q, k = self.rope(q, k, grid_hw=grid_hw) + attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=False, scale=self.scale) + attn_output = attn_output.transpose(1, 2).reshape( + bsz, seq_len, self.num_heads * self.head_dim + ) + return self.out_proj(attn_output) + + +class Step3VisionBlock(nn.Module): + """Single vision transformer block (Pre-LN + LayerScale).""" + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float, + hidden_act: str, + layer_norm_eps: float, + ls_init_value: Optional[float], + max_grid_height: int, + max_grid_width: int, + use_cls_token: bool, + use_rope2d: bool, + rope_theta: float, + rope_theta_rescale_factor: float, + ): + super().__init__() + self.attn = Step3VisionAttention( + hidden_size=hidden_size, + num_heads=num_heads, + max_grid_height=max_grid_height, + max_grid_width=max_grid_width, + use_cls_token=use_cls_token, + use_rope2d=use_rope2d, + rope_theta=rope_theta, + rope_theta_rescale_factor=rope_theta_rescale_factor, + ) + self.ln_1 = LayerNorm(hidden_size=hidden_size, eps=layer_norm_eps) + self.ln_2 = LayerNorm(hidden_size=hidden_size, eps=layer_norm_eps) + self.mlp = Step3VisionMLP(hidden_size, int(hidden_size * mlp_ratio), hidden_act) + ls = ls_init_value if ls_init_value is not None else 1.0 + self.ls_1 = Step3VisionLayerScale(hidden_size, ls) + self.ls_2 = Step3VisionLayerScale(hidden_size, ls) + + def forward( + self, + hidden_states: torch.Tensor, + grid_hw: Tuple[int, int], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + hidden_states = self.attn(hidden_states, grid_hw=grid_hw) + hidden_states = residual + self.ls_1(hidden_states) + + residual = hidden_states + hidden_states = self.ln_2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + self.ls_2(hidden_states) + return hidden_states + + +class Step3VisionTransformer(nn.Module): + def __init__(self, depth: int, **block_kwargs): + super().__init__() + self.resblocks = nn.ModuleList([Step3VisionBlock(**block_kwargs) for _ in range(depth)]) + + def forward( + self, + hidden_states: torch.Tensor, + grid_hw: Tuple[int, int], + ) -> torch.Tensor: + for block in self.resblocks: + hidden_states = block(hidden_states, grid_hw=grid_hw) + return hidden_states + + +class Step3p7VisionEncoder(nn.Module): + """Perception-Encoder vision tower (``vision_model.*`` checkpoint subtree). + + HF carries two trailing Conv2d downsamplers inside the same module + (``vit_downsampler1``, ``vit_downsampler2``) but invokes them externally in + ``Step3p7Model._process_image_features``. We keep the same parameter + layout so ``vision_model.vit_downsampler{1,2}.*`` weights load directly, + but call the downsamplers from ``forward`` here so a single call returns + the post-downsample feature map ready for the projector. + """ + + def __init__(self, vision_config: PretrainedConfig, dtype: torch.dtype): + super().__init__() + self.config = vision_config + + self.hidden_size = int(vision_config.width) + self.num_heads = int(vision_config.heads) + self.num_hidden_layers = int(vision_config.layers) + self.patch_size = int(vision_config.patch_size) + self.image_size = int(vision_config.image_size) + self.layer_norm_eps = float(getattr(vision_config, "layer_norm_eps", 1e-5)) + self.hidden_act = getattr(vision_config, "hidden_act", "quick_gelu") + self.mlp_ratio = float(getattr(vision_config, "mlp_ratio", 8960.0 / 1536.0)) + self.ls_init_value = getattr(vision_config, "ls_init_value", None) + self.use_cls_token = bool(getattr(vision_config, "use_cls_token", False)) + self.use_rope2d = bool(getattr(vision_config, "use_rope2d", True)) + self.use_abs_posemb = bool(getattr(vision_config, "use_abs_posemb", True)) + self.use_ln_pre = bool(getattr(vision_config, "use_ln_pre", True)) + self.use_ln_post = bool(getattr(vision_config, "use_ln_post", False)) + + self.conv1 = nn.Conv2d( + in_channels=int(getattr(vision_config, "num_channels", 3)), + out_channels=self.hidden_size, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) + self.ln_pre = ( + LayerNorm(hidden_size=self.hidden_size, eps=self.layer_norm_eps) + if self.use_ln_pre + else nn.Identity() + ) + self.ln_post = ( + LayerNorm(hidden_size=self.hidden_size, eps=self.layer_norm_eps) + if self.use_ln_post + else nn.Identity() + ) + + grid_size = self.image_size // self.patch_size + self.base_grid = (grid_size, grid_size) + + if self.use_cls_token: + self.class_embedding = nn.Parameter( + torch.randn(self.hidden_size) * (self.hidden_size**-0.5) + ) + else: + self.class_embedding = None + + if self.use_abs_posemb: + self.posemb_grid_size = grid_size + self.positional_embedding = nn.Parameter( + (self.hidden_size**-0.5) + * torch.randn( + int(self.use_cls_token) + self.posemb_grid_size**2, + self.hidden_size, + ) + ) + else: + self.posemb_grid_size = None + + self.transformer = Step3VisionTransformer( + depth=self.num_hidden_layers, + hidden_size=self.hidden_size, + num_heads=self.num_heads, + mlp_ratio=self.mlp_ratio, + hidden_act=self.hidden_act, + layer_norm_eps=self.layer_norm_eps, + ls_init_value=self.ls_init_value, + max_grid_height=grid_size, + max_grid_width=grid_size, + use_cls_token=self.use_cls_token, + use_rope2d=self.use_rope2d, + rope_theta=float(getattr(vision_config, "rope_theta", 10000.0)), + rope_theta_rescale_factor=float( + getattr(vision_config, "rope_theta_rescale_factor", 1.0) + ), + ) + + # Two trailing Conv2d downsamplers (stride 2 each, channel x2 each). + self.vit_downsampler1 = nn.Conv2d( + self.hidden_size, + self.hidden_size * 2, + kernel_size=3, + stride=2, + padding=1, + ) + self.vit_downsampler2 = nn.Conv2d( + self.hidden_size * 2, + self.hidden_size * 4, + kernel_size=3, + stride=2, + padding=1, + ) + + # Cast the whole tower to the model dtype (bf16 in practice). The + # PerceptionEncoder portion of the checkpoint is bf16 on disk even in + # the FP8 text checkpoint, so a single ``to(dtype)`` is sufficient. + self.to(dtype) + self._dtype = dtype + + def _sample_abs_posemb(self, grid_h: int, grid_w: int) -> torch.Tensor: + if self.posemb_grid_size == grid_h and self.posemb_grid_size == grid_w: + return self.positional_embedding[None, ...] + pos_embed = self.positional_embedding + if self.use_cls_token: + cls_token_embed, pos_embed = pos_embed[:1], pos_embed[1:] + pos_embed = ( + pos_embed.reshape(1, self.posemb_grid_size, self.posemb_grid_size, -1) + .permute(0, 3, 1, 2) + .contiguous() + ) + pos_embed = F.interpolate( + pos_embed, + size=(grid_h, grid_w), + mode="bilinear", + align_corners=False, + ) + pos_embed = pos_embed.permute(0, 2, 3, 1).reshape(-1, self.hidden_size) + if self.use_cls_token: + pos_embed = torch.cat([cls_token_embed, pos_embed], dim=0) + return pos_embed[None, ...] + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + bsz, _, height, width = pixel_values.shape + grid_h, grid_w = height // self.patch_size, width // self.patch_size + hidden = self.conv1(pixel_values) # (B, D, Gh, Gw) + hidden = hidden.flatten(2).transpose(1, 2) # (B, Gh*Gw, D) + if self.use_cls_token: + cls_token = self.class_embedding.view(1, 1, -1).expand(bsz, -1, -1) + hidden = torch.cat([cls_token, hidden], dim=1) + if self.use_abs_posemb: + hidden = hidden + self._sample_abs_posemb(grid_h, grid_w).to(hidden.dtype) + hidden = self.ln_pre(hidden) + hidden = self.transformer(hidden, grid_hw=(grid_h, grid_w)) + if self.use_ln_post: + hidden = self.ln_post(hidden) + if self.use_cls_token: + hidden = hidden[:, 1:, :] + return hidden + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Run vision transformer + the two trailing downsamplers. + + Returns features shaped ``(B, (Gh//4) * (Gw//4), 4 * width)`` ready + for the linear projector to map into the text hidden size. + """ + pixel_values = pixel_values.to(dtype=self._dtype, device=self.conv1.weight.device) + x = self.forward_features(pixel_values) # (B, P, D) + B, P, D = x.shape + T = int(P**0.5) + x = x.transpose(1, 2).contiguous().view(B, D, T, T) + x = self.vit_downsampler1(x) + x = self.vit_downsampler2(x) + Bd, Cd, Td, _ = x.shape + return x.view(Bd, Cd, Td * Td).transpose(1, 2) # (B, T'*T', 4D) + + +class Step3p7VisionTower(nn.Module): + """Vision encoder + projector exposed to TRT-LLM as the ``mm_encoder``. + + Encapsulates ``vision_model`` (``Step3p7VisionEncoder``) and the + ``vit_large_projector`` Linear from the HF checkpoint, and provides the + ``forward(multimodal_params)`` signature used by + ``get_multimodal_embeddings``. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]): + super().__init__() + pretrained_config = model_config.pretrained_config + vision_config = getattr(pretrained_config, "vision_config", None) + if vision_config is None: + raise ValueError( + "Step3p7VisionTower requires `vision_config` on the pretrained config; " + "the checkpoint is missing the vision tower entry." + ) + # Ensure vision_config carries a torch.dtype (not the raw JSON string). + _normalize_torch_dtype(vision_config) + # Honour the outer model dtype when set (typical bf16); fall back to + # the vision sub-config's stored dtype. + outer_dtype = getattr(pretrained_config, "torch_dtype", None) + if not isinstance(outer_dtype, torch.dtype): + outer_dtype = getattr(vision_config, "torch_dtype", None) + if not isinstance(outer_dtype, torch.dtype): + outer_dtype = torch.bfloat16 + self._dtype = outer_dtype + + text_config = _get_text_config(model_config) + self.image_token_id = int(getattr(pretrained_config, "image_token_id", 128001)) + + self.vision_model = Step3p7VisionEncoder(vision_config, dtype=outer_dtype) + + proj_in = 4 * int(vision_config.width) + proj_out = int(text_config.hidden_size) + proj_bias = bool(getattr(pretrained_config, "projector_bias", False)) + # Single GPU rank for the bring-up; keep the projector as a plain + # bf16 Linear so weights from ``vit_large_projector.weight`` load + # directly without an extra remapping. + self.vit_large_projector = nn.Linear(proj_in, proj_out, bias=proj_bias).to(outer_dtype) + + @property + def dtype(self) -> torch.dtype: + return self._dtype + + def load_weights(self, weights: Dict[str, torch.Tensor]): + """Consume ``vision_model.*`` and ``vit_large_projector.*`` keys. + + Both checkpoint layouts are supported: the FP8/BF16 exports key the + vision subtree directly (``vision_model.*`` / ``vit_large_projector.*``), + while the NVFP4 export nests everything under ``model.`` alongside the + ``model.language_model.*`` text decoder (``model.vision_model.*`` / + ``model.vit_large_projector.*``). An optional leading ``model.`` is + stripped before matching so both load identically. + + Returns silently if neither subtree is present (e.g. text-only + checkpoint slice); the caller will surface a missing-weights error + downstream via the model engine if that is wrong for the scenario. + """ + vision_state: Dict[str, torch.Tensor] = {} + projector_state: Dict[str, torch.Tensor] = {} + for key in list(weights.keys()): + # Normalize the NVFP4 ``model.`` wrapper prefix to the bare layout. + norm_key = key[len("model.") :] if key.startswith("model.") else key + if norm_key.startswith("vision_model."): + sub = norm_key[len("vision_model.") :] + vision_state[sub] = weights[key] + elif norm_key.startswith("vit_large_projector."): + sub = norm_key[len("vit_large_projector.") :] + projector_state[sub] = weights[key] + + if vision_state: + self.vision_model.load_state_dict(vision_state, strict=True) + if projector_state: + self.vit_large_projector.load_state_dict(projector_state, strict=True) + + def _process_image_features(self, image_features: torch.Tensor) -> torch.Tensor: + """Project post-downsample vision features to the text hidden size.""" + return self.vit_large_projector(image_features.to(self._dtype)) + + def _encode(self, pixel_values: torch.Tensor) -> torch.Tensor: + feats = self.vision_model(pixel_values) + return self._process_image_features(feats) + + @nvtx_range("Step3p7VisionTower forward()") + @torch.inference_mode() + def forward(self, multimodal_params: List[MultimodalParams]) -> List[torch.Tensor]: + """Encode all images (and their patches) in the batch. + + For each request we mirror the HF reference's per-image layout: + ``[patches_for_image_i ... | full_image_i]`` flattened into a single + ``(num_tokens, hidden)`` tensor. The caller (``fuse_input_embeds``) + concatenates per-request embeddings and writes them into the + positions where the placeholder token id appears in ``input_ids``. + + Token ordering must therefore match the input processor's text + placeholder expansion, which follows HF + ``Step3VLProcessor._get_image_repl_features``: patch features come + first (``num_patches`` blocks of ``num_patch_feature_size`` tokens), + then the full image feature block (``num_image_feature_size`` tokens). + """ + per_request_embeds: List[torch.Tensor] = [] + for mm in multimodal_params: + # TODO: batch full images across all requests, batch patch images across all + # requests where shapes match, and then split/reassemble according to num_patches + image_data = mm.multimodal_data.get("image") if mm.multimodal_data else None + if image_data is None: + continue + pixel_values = image_data.get("pixel_values") + patch_pixel_values = image_data.get("patch_pixel_values") + num_patches_list = image_data.get("num_patches") + if pixel_values is None: + continue + + pixel_values = pixel_values.to(self.vision_model.conv1.weight.device) + if pixel_values.dim() >= 5: + pixel_values = pixel_values.view(-1, *pixel_values.shape[-3:]) + elif pixel_values.dim() == 3: + pixel_values = pixel_values.unsqueeze(0) + image_feats = self._encode(pixel_values) # (N, P_img, H) + + patch_feats = None + if patch_pixel_values is not None and patch_pixel_values.numel() > 0: + patch_pixel_values = patch_pixel_values.to(self.vision_model.conv1.weight.device) + if patch_pixel_values.dim() >= 5: + patch_pixel_values = patch_pixel_values.view(-1, *patch_pixel_values.shape[-3:]) + elif patch_pixel_values.dim() == 3: + patch_pixel_values = patch_pixel_values.unsqueeze(0) + if patch_pixel_values.shape[0] > 0: + patch_feats = self._encode(patch_pixel_values) # (M, P_patch, H) + + # Build the per-request flat embedding stream: + # patches (if any, in order) then full image, repeated per image. + if num_patches_list is None: + num_images = image_feats.shape[0] + num_patches_list = [0] * num_images + elif isinstance(num_patches_list, torch.Tensor): + num_patches_list = num_patches_list.flatten().tolist() + else: + num_patches_list = list(num_patches_list) + + cur_patch_idx = 0 + flat_blocks: List[torch.Tensor] = [] + for img_idx, n_p in enumerate(num_patches_list): + if n_p > 0 and patch_feats is not None: + blk = patch_feats[cur_patch_idx : cur_patch_idx + n_p] + flat_blocks.append(blk.reshape(-1, blk.shape[-1])) + cur_patch_idx += n_p + flat_blocks.append(image_feats[img_idx].reshape(-1, image_feats.shape[-1])) + + if flat_blocks: + per_request_embeds.append(torch.cat(flat_blocks, dim=0)) + + if not per_request_embeds: + return [] + return [torch.cat(per_request_embeds, dim=0)] + + +# --------------------------------------------------------------------------- +# Multimodal input processor +# --------------------------------------------------------------------------- + + +class Step3p7VLInputProcessor(BaseMultimodalInputProcessor): + """Input processor wrapping the HF ``Step3VLProcessor`` (remote code). + + Produces a tokenized input stream with image-token placeholders rewritten + to an out-of-vocab sentinel (``vocab_size + 1``) so ``fuse_input_embeds`` + can locate them by ``input_ids >= vocab_size`` -- mirroring the + Qwen2.5-VL / LlavaNext pattern in TRT-LLM. The structural special tokens + (````, ````, ````, ````, + ````) keep their original token ids; the language model + embeds those normally. + """ + + def __init__( + self, + model_path: str, + config: PretrainedConfig, + tokenizer: Optional[AutoTokenizer], + trust_remote_code: bool = True, + **kwargs, + ): + super().__init__( + model_path=model_path, + config=config, + tokenizer=tokenizer, + trust_remote_code=trust_remote_code, + **kwargs, + ) + self._tokenizer = ( + tokenizer + if tokenizer is not None + else AutoTokenizer.from_pretrained(model_path, trust_remote_code=trust_remote_code) + ) + # The Step3p7 checkpoint ships ``processing_step3.Step3VLProcessor`` as + # a remote-code module, but does not register an ``AutoProcessor`` entry + # nor a ``processor_config.json``. ``AutoProcessor.from_pretrained`` + # therefore falls back to a tokenizer-only path that silently drops + # the image input. Load the remote class directly so images are + # actually preprocessed. + # + # ``Step3VLProcessor`` expects a raw HF tokenizer (uses ``get_vocab``). + # The runtime may hand us a ``TransformersTokenizer`` wrapper instead; + # unwrap it before constructing the processor. + hf_tokenizer = getattr(self._tokenizer, "tokenizer", self._tokenizer) + try: + processor_cls = get_class_from_dynamic_module( + "processing_step3.Step3VLProcessor", + model_path, + ) + self._processor = processor_cls( + tokenizer=hf_tokenizer, + chat_template=getattr(hf_tokenizer, "chat_template", None), + ) + except Exception: + logger.warning( + "[Step3p7VL] Falling back to AutoProcessor; image inputs may not be processed." + ) + self._processor = AutoProcessor.from_pretrained( + model_path, + use_fast=self.use_fast, + trust_remote_code=trust_remote_code, + ) + # The remote ``Step3VLProcessor`` exposes ``image_token_id`` via the + # tokenizer vocabulary, but its getter calls ``get_vocab()`` which + # raises ``NotImplementedError`` on tokenizers loaded without the + # fast backend. Prefer the model config value, falling back to the + # processor only when it is safely accessible. + cfg_image_token_id = getattr(config, "image_token_id", None) + if cfg_image_token_id is None: + try: + cfg_image_token_id = self._processor.image_token_id # type: ignore[attr-defined] + except Exception: + cfg_image_token_id = 128001 + self._image_token_id = int(cfg_image_token_id) + text_config = getattr(config, "text_config", config) + self._dtype = getattr(text_config, "torch_dtype", torch.bfloat16) + if isinstance(self._dtype, str): + self._dtype = getattr(torch, self._dtype, torch.bfloat16) + self._vocab_size = int(getattr(text_config, "vocab_size", 0)) + self._tllm_multimodal_token_id = self._vocab_size + 1 + + hf_tok = getattr(self._tokenizer, "tokenizer", self._tokenizer) + special_token_ids: List[int] = [] + for tok in ("", "", "", "", ""): + try: + tok_id = hf_tok.convert_tokens_to_ids(tok) + except Exception: + tok_id = None + unk_id = getattr(hf_tok, "unk_token_id", None) + if tok_id is None or (unk_id is not None and tok_id == unk_id): + logger.warning( + "[Step3p7VL] Could not resolve structural token %r; " + "multimodal hashing will fall back to the vocab-size " + "discriminator only.", + tok, + ) + special_token_ids = [] + break + special_token_ids.append(int(tok_id)) + self._mm_special_token_ids = special_token_ids + + # ------- BaseMultimodalInputProcessor required properties ------------- + + @property + def processor(self) -> AutoProcessor: + return self._processor + + @property + def tokenizer(self) -> AutoTokenizer: + return self._tokenizer + + @property + def config(self) -> PretrainedConfig: + return self._config + + @property + def dtype(self) -> torch.dtype: + return self._dtype + + # ---------------------------------------------------------------------- + + def get_vocab_size(self) -> int: + return self._vocab_size + + def get_num_tokens_per_image( + self, + *, + image: Union[Image.Image, torch.Tensor], + **kwargs, + ) -> int: + """Total prompt tokens for one image, framing tokens included. + + Delegates to the remote processor's ``get_num_image_tokens`` (which + accounts for ````/````/```` per + tile and ````/```` for the global features), so the + count matches the contiguous span produced by ``call_with_text_prompt``. + """ + if isinstance(image, torch.Tensor): + height, width = int(image.shape[-2]), int(image.shape[-1]) + else: + width, height = image.width, image.height + return int(self._processor.get_num_image_tokens(width, height)) + + def get_mm_token_ids(self) -> Optional[torch.Tensor]: + """Token ids forming one logical image unit (embed slots + framing). + + ``call_with_text_prompt`` rewrites the ```` placeholders to the OOV + sentinel ``vocab_size + 1``, so the embed slots are matched by that + sentinel rather than the original ```` id. Returning the + sentinel together with the framing tokens keeps the whole image span + contiguous in the hashing mask. Falls back to ``None`` (vocab-size + discriminator) when the framing tokens could not be resolved. + """ + if not self._mm_special_token_ids: + return None + return torch.tensor([self._tllm_multimodal_token_id] + self._mm_special_token_ids) + + def get_mm_special_token_ids(self) -> Optional[torch.Tensor]: + """Framing-token ids inside an image span that carry no vision embed. + + These are subtracted from the embed-row mask so the embed-slot count + stays accurate while the span itself remains contiguous. + """ + if not self._mm_special_token_ids: + return None + return torch.tensor(self._mm_special_token_ids) + + @torch.inference_mode() + def call_with_text_prompt( + self, inputs: TextPrompt, sampling_params: SamplingParams + ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: + text_prompt = inputs.get("prompt") + mm_data = inputs.get("multi_modal_data") or {} + images = mm_data.get("image", []) if isinstance(mm_data, dict) else [] + + if not images: + token_ids = self._tokenizer(text_prompt, return_tensors="pt").input_ids[0] + return token_ids.to(torch.int32).tolist(), {} + + processed = self._processor( + text=text_prompt, + images=images, + return_tensors="pt", + ) + input_ids = processed["input_ids"][0] + + # Replace each position with the OOV sentinel; structural + # tokens (, , etc.) keep their normal ids. + sentinel = self._tllm_multimodal_token_id + input_ids = input_ids.clone() + input_ids[input_ids == self._image_token_id] = sentinel + + multimodal_data: Dict[str, Any] = {"image": {}} + image_dict = multimodal_data["image"] + image_dict["pixel_values"] = processed["pixel_values"].to(self._dtype) + if "patch_pixel_values" in processed: + image_dict["patch_pixel_values"] = processed["patch_pixel_values"].to(self._dtype) + if "num_patches" in processed: + np_val = processed["num_patches"] + image_dict["num_patches"] = ( + np_val.tolist() if isinstance(np_val, torch.Tensor) else list(np_val) + ) + if "patch_newline_mask" in processed: + image_dict["patch_newline_mask"] = processed["patch_newline_mask"] + + return input_ids.to(torch.int32).tolist(), {"multimodal_data": multimodal_data} + + +# --------------------------------------------------------------------------- +# VLM wrapper (vision + language) +# --------------------------------------------------------------------------- + + +@register_vision_encoder(Step3p7VisionTower) +@register_auto_model("Step3p7ForConditionalGeneration") +@register_input_processor( + Step3p7VLInputProcessor, + model_type="step3p7", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={"image": ""}, + content_format=ContentFormat.OPENAI, + ), +) +class Step3p7VLForConditionalGeneration(nn.Module): + """Multimodal entry point for the Step3p7 Flash checkpoint. + + Wraps the existing text-only ``Step3p7ForCausalLM`` plus the + PerceptionEncoder-based vision tower. When a request has no + multimodal payload this class is a thin passthrough so plain text + benchmarks continue to use the original Step3p7 text path. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs): + super().__init__() + self.model_config = model_config + self.config = model_config.pretrained_config + + # Mirror text-config aliases up onto the top-level config (mirrors + # what Step3p7ForCausalLM.__init__ does) before instantiating the + # text model; otherwise the inner __init__ would be working on a + # config whose torch_dtype etc. are still strings. + _mirror_step3p7_text_aliases(self.config) + + # Inner causal LM: reuses all the existing decoder + MTP wiring. + self.llm = Step3p7ForCausalLM(model_config) + + # Vision encoder. Built lazily in load_weights() (outside MetaInitMode) + # so its PerceptionEncoder / HF submodules allocate real tensors instead + # of meta tensors. This keeps the large text LLM on the fast meta-init + # path while avoiding leftover meta tensors that would crash at runtime. + # Stays None in the disaggregated decode worker. + self.mm_encoder = None + + # ----- engine-facing surface (delegate to the inner causal LM) ------- + + @property + def vocab_size_padded(self) -> int: + return self.llm.vocab_size_padded + + def infer_max_seq_len(self) -> int: + return self.llm.infer_max_seq_len() + + @property + def model(self): + # Expose the inner decoder model so existing helpers that walk + # ``self.model.layers`` (e.g. weight bookkeeping) keep working when + # they see the wrapper instead of the bare causal LM. + return self.llm.model + + @property + def lm_head(self): + return self.llm.lm_head + + @property + def epilogue(self): + return self.llm.epilogue + + @property + def spec_worker(self): + return getattr(self.llm, "spec_worker", None) + + @property + def draft_model(self): + return getattr(self.llm, "draft_model", None) + + # --------------------------------------------------------------------- + + def load_weights( + self, + weights, + weight_mapper=None, + skip_modules=None, + params_map=None, + allow_partial_loading: bool = False, + ): + """Split vision/text weights and delegate to the inner LM loader.""" + if self.mm_encoder is None and not _is_disagg() and hasattr(weights, "items"): + # Construct the vision tower here, outside MetaInitMode, so its + # PerceptionEncoder / HF submodules allocate real tensors. Move it + # straight to CUDA (model_loader already ran model.to("cuda") for + # the LLM); the load_state_dict below copies the checkpoint in. + self.mm_encoder = Step3p7VisionTower(self.model_config).eval().to("cuda") + if self.mm_encoder is not None and hasattr(weights, "items"): + # Hand the vision subtree to the encoder; it consumes the keys + # in-place via state_dict (no removal from ``weights`` needed — + # ``Step3p7ForCausalLM.load_weights`` already strips + # ``vision_model.`` / ``vit_large_projector.`` via + # ``ignored_key_prefixes``). + self.mm_encoder.load_weights(weights) + return self.llm.load_weights( + weights, + weight_mapper=weight_mapper, + skip_modules=skip_modules, + params_map=params_map, + allow_partial_loading=allow_partial_loading, + ) + + @torch.inference_mode() + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + multimodal_params = kwargs.pop("multimodal_params", []) + num_context_requests = attn_metadata.num_contexts + + mm_embeds: List[torch.Tensor] = [] + # Only context requests carry pixel values; generation steps after + # the first iteration have no image payload to encode. + mm_context_params = [ + p + for p in multimodal_params[:num_context_requests] + if ( + p.multimodal_data is not None + and ( + p.multimodal_data.get("image", {}).get("pixel_values") is not None + or p.multimodal_data.get("multimodal_embedding") is not None + ) + ) + ] + if mm_context_params and self.mm_encoder is not None: + mm_embeds = get_multimodal_embeddings( + encoder_forward_fn=self.mm_encoder.forward, + multimodal_params=mm_context_params, + ) + mm_embeds = find_input_mm_embeds(mm_embeds, mm_context_params) + + if input_ids is not None and mm_embeds: + input_ids, inputs_embeds = fuse_input_embeds( + self.llm.model.embed_tokens, input_ids, mm_embeds, **kwargs + ) + + return self.llm.forward( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + return_context_logits=return_context_logits, + spec_metadata=spec_metadata, + resource_manager=resource_manager, + **kwargs, + ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py b/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py index 8376f77f497e..503106033ea6 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py @@ -538,6 +538,8 @@ def cvt_routing_method_type(self, routing_method_type) -> int: _trtllm.DeepSeekV3: _flashinfer.DeepSeekV3, _trtllm.Llama4: _flashinfer.Llama4, _trtllm.RenormalizeNaive: _flashinfer.RenormalizeNaive, + _trtllm.MiniMax2: _flashinfer.MiniMax2, + _trtllm.SigmoidRenorm: _flashinfer.SigmoidRenorm, _trtllm.Unspecified: _flashinfer.Unspecified, } if routing_method_type not in _mapping: diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index bfaf761c1b05..ef6910479bbc 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -299,6 +299,13 @@ def load_config_and_apply_defaults( config = checkpoint_loader.load_config(checkpoint_dir, **config_kwargs) + if llm_args.speculative_config is not None: + from tensorrt_llm._torch.speculative import \ + update_spec_config_from_model_config + + update_spec_config_from_model_config(llm_args.speculative_config, + config.pretrained_config) + model_cls = AutoModelForCausalLM._resolve_class(config) # model_cls is None when the architecture is unknown/unsupported. diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 20181bf29179..b8b5ff102586 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -892,6 +892,17 @@ def change_attn_metadata(self, num_accepted_tokens: torch.Tensor, attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( mtp_num_modules + 1 - num_accepted_tokens[num_contexts:batch_size]) + # A generation request's KV length can never be smaller than the + # number of query tokens (mtp_num_modules) the draft layer processes + # this step. This only underflows for the tiny dummy sequences used + # during generation-step warmup; real sequences are long enough that + # subtracting the rejected draft tokens stays well above the draft + # length. Clamping in-place (CUDA-graph safe) keeps an invalid + # kv_len < q_len from reaching the attention kernel, which otherwise + # computes a negative effective KV range and triggers an illegal + # memory access (e.g. Step3p7 MTP with dense sliding-window attention). + attn_metadata.kv_lens_cuda[num_contexts:batch_size].clamp_( + min=mtp_num_modules) attn_metadata.on_update_kv_lens() if attn_metadata.kv_cache_params is not None and not attn_metadata.is_cuda_graph: diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 68ef347b0533..f2dd8ed2ca50 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -381,29 +381,39 @@ def update_spec_config_from_model_config(spec_config, model_config): from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig if not isinstance(spec_config, MTPDecodingConfig): return - # Read num_nextn_predict_layers from the model's pretrained config. - # This determines the actual MTP layer count in the checkpoint and drives - # the spec_dec_mode decision (EAGLE vs vanilla MTP). - spec_config.num_nextn_predict_layers = model_config.num_nextn_predict_layers - # For vanilla MTP (>1 MTP layers in the checkpoint): set max_draft_len to the - # minimum of the user-requested value and the model's layer count. - # If the user explicitly requested fewer draft tokens than the model has layers, - # honour that and warn. Otherwise default to using all model layers. - if spec_config.spec_dec_mode.is_mtp_vanilla(): - model_layers = spec_config.num_nextn_predict_layers - user_set = 'max_draft_len' in spec_config.model_fields_set - if user_set and spec_config.max_draft_len < model_layers: - logger.warning( - f"MTP: max_draft_len ({spec_config.max_draft_len}) is less than the model's " - f"num_nextn_predict_layers ({model_layers}). " - f"Using max_draft_len={spec_config.max_draft_len} draft tokens." - ) - # Keep the user-set max_draft_len as-is - else: - spec_config.max_draft_len = model_layers - spec_config.max_total_draft_tokens = spec_config.max_draft_len - # For Eagle MTP (1 MTP layer): max_draft_len controls how many times the - # single layer is run. It was already set by the user (defaults to 1). + # Read the MTP layer count from the model's pretrained config. This + # determines the actual MTP layer count in the checkpoint and drives the + # spec_dec_mode decision (EAGLE vs vanilla MTP). Different checkpoints expose + # this under different names: DeepSeek-style configs use + # `num_nextn_predict_layers`, while Qwen3Next-style configs (including + # Qwen3.5) use `mtp_num_hidden_layers`. Fall back to a single shared MTP / + # EAGLE layer when neither field is present. + num_nextn_predict_layers = getattr(model_config, "num_nextn_predict_layers", + None) + if num_nextn_predict_layers is None: + num_nextn_predict_layers = getattr(model_config, + "mtp_num_hidden_layers", None) + if num_nextn_predict_layers is None: + num_nextn_predict_layers = 1 + spec_config.num_nextn_predict_layers = num_nextn_predict_layers + is_vanilla = spec_config.spec_dec_mode.is_mtp_vanilla() + + # Resolve max_draft_len when the user didn't set it: + # vanilla MTP -> use all checkpoint MTP heads + # MTP-Eagle -> replay the single head once + if spec_config.max_draft_len is None: + spec_config.max_draft_len = (spec_config.num_nextn_predict_layers + if is_vanilla else 1) + elif is_vanilla and spec_config.max_draft_len != spec_config.num_nextn_predict_layers: + effective_draft_len = min(spec_config.max_draft_len, + spec_config.num_nextn_predict_layers) + logger.warning( + f"MTP: max_draft_len ({spec_config.max_draft_len}) does not match " + f"num_nextn_predict_layers ({spec_config.num_nextn_predict_layers}); " + f"using max_draft_len={effective_draft_len} draft tokens.") + spec_config.max_draft_len = effective_draft_len + + spec_config.max_total_draft_tokens = spec_config.max_draft_len @dataclass diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index 2b9d810c34ec..d41ede2e4180 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -665,20 +665,34 @@ def evaluate(self, def command_harness(cls, ctx, **kwargs): llm: Union[LLM, PyTorchLLM] = ctx.obj - evaluator = cls(dataset_path=kwargs.pop("dataset_path", None), - num_samples=kwargs.pop("num_samples", None), - random_seed=kwargs.pop("random_seed", 0), - apply_chat_template=kwargs.pop("apply_chat_template", - False), - fewshot_as_multiturn=kwargs.pop("fewshot_as_multiturn", - False), - system_prompt=kwargs.pop("system_prompt", None), - is_multimodal=kwargs.pop("is_multimodal", False), - chat_template_kwargs=kwargs.pop("chat_template_kwargs", - None), - log_samples=kwargs.pop("log_samples", False), - output_path=kwargs.pop("output_path", None), - output_dir=kwargs.pop("output_dir", None)) + # Resolve the post-processor: accept a callable (already-bound) or the + # string key "strip_thinking_mmmu" coming from CLI flags. + post_process_fn = kwargs.pop("post_process_fn", None) + if isinstance(post_process_fn, str): + if post_process_fn == "strip_thinking_mmmu": + from .post_processing import \ + strip_thinking_and_extract_mmmu_answer + post_process_fn = strip_thinking_and_extract_mmmu_answer + else: + raise click.BadParameter( + f"Unknown --post_process_fn={post_process_fn!r}; expected 'strip_thinking_mmmu'." + ) + + evaluator = cls( + dataset_path=kwargs.pop("dataset_path", None), + num_samples=kwargs.pop("num_samples", None), + random_seed=kwargs.pop("random_seed", 0), + apply_chat_template=kwargs.pop("apply_chat_template", False), + fewshot_as_multiturn=kwargs.pop("fewshot_as_multiturn", False), + system_prompt=kwargs.pop("system_prompt", None), + is_multimodal=kwargs.pop("is_multimodal", False), + chat_template_kwargs=kwargs.pop("chat_template_kwargs", None), + log_samples=kwargs.pop("log_samples", False), + output_path=kwargs.pop("output_path", None), + output_dir=kwargs.pop("output_dir", None), + post_process_fn=post_process_fn, + preserve_caller_max_tokens=kwargs.pop("preserve_caller_max_tokens", + False)) # Optional sampling overrides (default: greedy, as before). # When any of temperature / top_p / top_k / seed is set, the wrapper # uses CLI values in preference to the task yaml's gen_kwargs so @@ -1252,6 +1266,21 @@ def __init__(self, **kwargs): type=str, default=None, help="Directory to save the task infos.") + @click.option( + "--preserve_caller_max_tokens", + is_flag=True, + default=False, + help="Keep --max_output_length when larger than lm-eval task default " + "(MMMU's max_gen_toks=512 is too small for thinking models that " + "produce chain-of-thought before the answer).") + @click.option( + "--post_process_fn", + type=click.Choice(["strip_thinking_mmmu"]), + default=None, + help="Per-sample post-processor. 'strip_thinking_mmmu' strips " + "... and then runs the MMMU answer extractor — needed " + "for thinking models (Kimi K2.5, Step3p7) whose CoT output the " + "default lm-eval regex cannot parse.") @click.pass_context @staticmethod def command(ctx, **kwargs) -> None: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 9e7739ad104a..167f4b443873 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1818,14 +1818,15 @@ def _remap_deprecated_num_nextn_predict_layers(cls, data): @model_validator(mode="after") def set_max_total_draft_tokens(self): - # Default max_draft_len to 1 if not set by the user. - # For vanilla MTP, update_spec_config_from_model_config will override this - # with the actual num_nextn_predict_layers from the model. - if self.max_draft_len is None: - self.max_draft_len = 1 - elif self.max_draft_len <= 0: - raise ValueError("max_draft_len must be > 0 for MTP") - self.max_total_draft_tokens = self.max_draft_len # Current MTP only supports linear tree + # Leave max_draft_len as None ("use the model's num_nextn_predict_layers") + # when the user doesn't set it; update_spec_config_from_model_config + # resolves it from the checkpoint before the model runs. When the user + # does set it, validate and mirror to max_total_draft_tokens (current MTP + # only supports a linear tree). + if self.max_draft_len is not None: + if self.max_draft_len <= 0: + raise ValueError("max_draft_len must be > 0 for MTP") + self.max_total_draft_tokens = self.max_draft_len return self @model_validator(mode="after") diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 2b2b4e1fc43e..769a0f8bc65e 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -470,3 +470,17 @@ zai-org/GLM-5-FP8: - quant_algo: FP8_BLOCK_SCALES spec_dec_algo: MTP accuracy: 78.0 +# Step-3.7-Flash text decoder (MoE) GSM8K, full 1319-sample split, TP4/EP4 with +# TRTLLM attention + MoE backends. FP8 measured on the FP8 block-scale +# checkpoint; NVFP4 measured on the modelopt NVFP4 export (FP8 KV cache). FP8 +# MTP (mtp_nextn=3) is lossless and reuses the non-spec baseline. The NVFP4 +# export does not ship MTP weights, so it is only graded without MTP. +stepfun-ai/Step-3.7-Flash: + - quant_algo: FP8_BLOCK_SCALES + accuracy: 88 + - quant_algo: FP8_BLOCK_SCALES + spec_dec_algo: MTP + accuracy: 88 + - quant_algo: NVFP4 + kv_cache_quant_algo: FP8 + accuracy: 88 diff --git a/tests/integration/defs/accuracy/references/mmmu.yaml b/tests/integration/defs/accuracy/references/mmmu.yaml index 883f6b63410a..5e3a2ba8f4c9 100644 --- a/tests/integration/defs/accuracy/references/mmmu.yaml +++ b/tests/integration/defs/accuracy/references/mmmu.yaml @@ -77,3 +77,15 @@ moonshotai/Kimi-K2.5: - quant_algo: NVFP4 kv_cache_quant_algo: FP8 accuracy: 81.56 +# Step-3.7-Flash multimodal (PerceptionEncoder vision tower + MoE text decoder). +# Reasoning model: ... traces are stripped before MMMU answer +# extraction (see test_llm_api_pytorch_multimodal.py::TestStep3_7). FP8 measured +# on the FP8 block-scale checkpoint over the full 900-sample MMMU val split (TP4 +# / EP4, TRTLLM attention + MoE backends); NVFP4 (modelopt export, FP8 KV cache) +# reuses the FP8 baseline. +stepfun-ai/Step-3.7-Flash: + - quant_algo: FP8_BLOCK_SCALES + accuracy: 64 + - quant_algo: NVFP4 + kv_cache_quant_algo: FP8 + accuracy: 64 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 6ca44610dcd8..a56816ea4832 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7423,3 +7423,84 @@ def test_8gpus(self, tp_size, ep_size): **pytorch_config) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + + +class TestStep3_7(LlmapiAccuracyTestHarness): + # Step-3.7-Flash is a MoE model registered under the multimodal + # architecture (Step3p7ForConditionalGeneration); text-only GSM8K exercises + # the text decoder path. The custom HF config requires trust_remote_code. + MODEL_NAME = "stepfun-ai/Step-3.7-Flash" + + @pytest.mark.skip_less_device(8) + @pytest.mark.skip_less_device_memory(140000) + @parametrize_with_ids("tp_size,ep_size", [(8, 8)]) + def test_auto_dtype(self, tp_size, ep_size): + model_path = f"{llm_models_root()}/Step-3.7-Flash" + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7) + with LLM(model_path, + tensor_parallel_size=tp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + max_seq_len=8192, + trust_remote_code=True) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(80000) + @parametrize_with_ids("mtp_nextn", [0, 3]) + @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) + def test_fp8_block_scales(self, tp_size, ep_size, mtp_nextn): + model_path = f"{llm_models_root()}/Step-3.7-Flash-FP8" + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7, + use_kv_cache_manager_v2=True) + pytorch_config = dict( + disable_overlap_scheduler=False, + cuda_graph_config=CudaGraphConfig(enable_padding=True), + moe_config=MoeConfig(backend="TRTLLM"), + ) + + mtp_config = None + if mtp_nextn > 0: + mtp_config = MTPDecodingConfig(max_draft_len=mtp_nextn) + + with LLM(model_path, + tensor_parallel_size=tp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + max_seq_len=8192, + attn_backend="TRTLLM", + speculative_config=mtp_config, + trust_remote_code=True, + **pytorch_config) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + + @skip_pre_blackwell + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(80000) + @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) + def test_nvfp4(self, tp_size, ep_size): + # The NVFP4 export does not ship MTP weights, so this checkpoint is only + # exercised without speculative decoding (unlike the FP8/BF16 ones). + model_path = f"{llm_models_root()}/Step-3.7-Flash-NVFP4" + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7, + use_kv_cache_manager_v2=True) + pytorch_config = dict( + disable_overlap_scheduler=False, + cuda_graph_config=CudaGraphConfig(enable_padding=True), + moe_config=MoeConfig(backend="TRTLLM"), + ) + + with LLM(model_path, + tensor_parallel_size=tp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + max_seq_len=8192, + attn_backend="TRTLLM", + trust_remote_code=True, + **pytorch_config) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py b/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py index c4a1c29f44d6..7a84a910fc9a 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py @@ -738,3 +738,74 @@ def test_auto_dtype( sampling_params=sampling_params, extra_evaluator_kwargs=extra_evaluator_kwargs, ) + + +class TestStep3_7(LlmapiAccuracyTestHarness): + # Step-3.7-Flash is a reasoning VLM: a PerceptionEncoder vision tower plus a + # MoE text decoder, registered under the Step3p7ForConditionalGeneration + # architecture (custom HF config -> trust_remote_code). MMMU exercises the + # vision path end to end. The model emits ... traces before + # its answer, so we strip them and extract the final MMMU letter (same + # handling as Kimi K2.5); preserve_caller_max_tokens keeps our larger + # generation budget instead of lm-eval's 512-token default (too small for + # the chain-of-thought). The text-only GSM8K path lives in + # test_llm_api_pytorch.py::TestStep3_7. + MODEL_NAME = "stepfun-ai/Step-3.7-Flash" + + # Validated with --max_input_length / --max_output_length 4096. + sampling_params = SamplingParams( + max_tokens=4096, + truncate_prompt_tokens=4096, + ) + + EXTRA_EVALUATOR_KWARGS = dict( + post_process_fn=strip_thinking_and_extract_mmmu_answer, + preserve_caller_max_tokens=True, + ) + + kv_cache_config = KvCacheConfig( + free_gpu_memory_fraction=0.7, + use_kv_cache_manager_v2=True, + ) + + def _make_llm(self, model_path: str): + pytorch_config = dict( + disable_overlap_scheduler=False, + cuda_graph_config=CudaGraphConfig(enable_padding=False), + moe_config=MoeConfig(backend="TRTLLM"), + ) + return LLM( + model_path, + tensor_parallel_size=4, + moe_expert_parallel_size=4, + kv_cache_config=self.kv_cache_config, + max_seq_len=8192, + attn_backend="TRTLLM", + trust_remote_code=True, + **pytorch_config, + ) + + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(80000) + def test_fp8_block_scales(self): + with self._make_llm(f"{llm_models_root()}/Step-3.7-Flash-FP8") as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + task = MMMU(self.MODEL_NAME) + task.evaluate( + llm, + sampling_params=self.sampling_params, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS, + ) + + @skip_pre_blackwell + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(80000) + def test_nvfp4(self): + with self._make_llm(f"{llm_models_root()}/Step-3.7-Flash-NVFP4") as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 + task = MMMU(self.MODEL_NAME) + task.evaluate( + llm, + sampling_params=self.sampling_params, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS, + ) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 730725662353..5526145cae72 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -776,6 +776,10 @@ accuracy/test_llm_api_pytorch.py::TestQwen3NextInstruct::test_nvfp4[tp4ep4_adp_o accuracy/test_llm_api_pytorch.py::TestQwen3NextInstruct::test_nvfp4[tp4ep4_adp_on-trtllm] accuracy/test_llm_api_pytorch.py::TestQwen3NextThinking::test_auto_dtype[tp4ep4] accuracy/test_llm_api_pytorch.py::TestSeedOss_36B::test_auto_dtype +accuracy/test_llm_api_pytorch.py::TestStep3_7::test_auto_dtype[tp_size=8-ep_size=8] TIMEOUT (90) +accuracy/test_llm_api_pytorch.py::TestStep3_7::test_fp8_block_scales[tp_size=4-ep_size=4-mtp_nextn=0] TIMEOUT (90) +accuracy/test_llm_api_pytorch.py::TestStep3_7::test_fp8_block_scales[tp_size=4-ep_size=4-mtp_nextn=3] TIMEOUT (90) +accuracy/test_llm_api_pytorch.py::TestStep3_7::test_nvfp4[tp_size=4-ep_size=4] TIMEOUT (90) accuracy/test_llm_api_pytorch_encode.py::TestEncoderEncode::test_encoder_encode_matches_huggingface_classification[bert-yelp-eager] accuracy/test_llm_api_pytorch_encode.py::TestEncoderEncode::test_encoder_encode_matches_huggingface_classification[bert-yelp-cuda_graph] accuracy/test_llm_api_pytorch_encode.py::TestEncoderEncode::test_encoder_encode_cuda_graph_matches_eager_logits[bert-yelp] @@ -802,6 +806,8 @@ accuracy/test_llm_api_pytorch_multimodal.py::TestQwen2_VL_7B::test_auto_dtype accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3VL::test_auto_dtype[forced_chunked_prefill] accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3VL_MOE::test_auto_dtype accuracy/test_llm_api_pytorch_multimodal.py::TestKimiK25::test_nvfp4[dep8] +accuracy/test_llm_api_pytorch_multimodal.py::TestStep3_7::test_fp8_block_scales TIMEOUT (120) +accuracy/test_llm_api_pytorch_multimodal.py::TestStep3_7::test_nvfp4 TIMEOUT (120) accuracy/test_llm_api_pytorch_multimodal.py::TestVILA1_5_3B::test_auto_dtype accuracy/test_llm_api_pytorch_ray.py::TestLlama3_1_8BInstruct::test_pp2_ray unittest/disaggregated/test_openai_disagg_server.py diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index ef235f800c9b..7007503f3e1b 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -26,6 +26,7 @@ l0_a10: - unittest/_torch/modeling/test_nemotron_nano_preprocessing.py - unittest/_torch/modeling/test_modeling_parakeet.py - unittest/_torch/modeling/test_modeling_radio.py + - unittest/_torch/modeling/test_modeling_step3p7.py - unittest/_torch/sampler/test_trtllm_sampler.py - unittest/_torch/executor/test_async_transfer_manager.py - unittest/_torch/executor/test_scheduler_serializable_output.py diff --git a/tests/integration/test_lists/test-db/l0_l40s.yml b/tests/integration/test_lists/test-db/l0_l40s.yml index 1817442a7ab4..b76f207dd9c1 100644 --- a/tests/integration/test_lists/test-db/l0_l40s.yml +++ b/tests/integration/test_lists/test-db/l0_l40s.yml @@ -24,6 +24,7 @@ l0_l40s: - unittest/_torch/modeling/test_modeling_qwen3vl_moe.py::TestQwen3VLMoe::test_all - unittest/_torch/modeling/test_modeling_qwen3vl.py::TestQwen3VL::test_all - unittest/_torch/modeling/test_modeling_qwen3vl.py::test_qwen3vl_init_preserves_caller_quant_config + - unittest/_torch/modeling/test_modeling_step3p7vl.py - test_e2e.py::test_ptp_scaffolding[DeepSeek-R1-Distill-Qwen-7B-DeepSeek-R1/DeepSeek-R1-Distill-Qwen-7B] - unittest/llmapi/apps/_test_openai_chat_multimodal.py::test_single_chat_session_image_embeds -m needs_l40s # MMMU sanity check diff --git a/tests/unittest/_torch/modeling/test_modeling_step3p7.py b/tests/unittest/_torch/modeling/test_modeling_step3p7.py new file mode 100644 index 000000000000..8f4b9f5c5a96 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_modeling_step3p7.py @@ -0,0 +1,715 @@ +# 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. +"""Tests for the Step3p7 text-generation bring-up (Step3p7ForConditionalGeneration). + +TestStep3p7Helpers — pure-Python loader/helper tests (no checkpoint, no GPU): + - Stacked routed-expert weight splitting (FP8 block-scale and NVFP4 layouts) + - MTP weight rewriting to the ``mtp_block`` layout + - ``model.language_model.*`` namespace flattening for the multimodal checkpoint + - quant-config exclude-module normalization + - NVFP4 dequant round-tripping + - MTP shared-head normalization before the output projection + +TestStep3p7AutoModelRegistration — model registry verification: + - ``Step3p7ForConditionalGeneration`` resolves to the VLM entry point + - ``Step3p5ForCausalLM`` resolves to the text-only causal LM + +TestStep3p7Checkpoint — checkpoint-backed config / weight-accounting tests +(requires the Step-3.7-Flash checkpoints under ``LLM_MODELS_ROOT``): + - Config sanity, per-layer attention/MoE inventory, and FP8 quant config + - Weight accounting that separates consumed text-path keys from intentionally + ignored vision and plain-path MTP keys + - Per-layer head / RoPE / SwiGLU helpers against the real config + - Default ``MTPDecodingConfig`` resolving to the checkpoint layer count +""" + +import json +import os +import types +import unittest + +import torch +from parameterized import parameterized +from transformers import PretrainedConfig +from utils.llm_data import llm_models_root + +# Resolve the Step3p7 checkpoints under the shared model root (LLM_MODELS_ROOT) +# like the other modeling tests instead of hard-coding a developer workspace +# path. The FP8 block-scale, NVFP4, and BF16 reference checkpoints share the +# same per-layer geometry; only the routed-expert dtype/layout differs. +# +# The FP8 and BF16 checkpoints store text decoder keys directly under +# ``model.layers.*`` and ship 3 plain-path MTP layers (45..47). The NVFP4 +# checkpoint is a modelopt export: text keys live under +# ``model.language_model.layers.*``, vision keys under ``model.vision_model.*``, +# and it carries no MTP layers. +STEP3P7_FP8_DIR = str(os.path.join(llm_models_root(), "Step-3.7-Flash-FP8")) +STEP3P7_NVFP4_DIR = str(os.path.join(llm_models_root(), "Step-3.7-Flash-NVFP4")) +STEP3P7_BF16_DIR = str(os.path.join(llm_models_root(), "Step-3.7-Flash")) + + +def _load_config(checkpoint_dir: str) -> dict: + with open(os.path.join(checkpoint_dir, "config.json")) as f: + return json.load(f) + + +def _load_safetensors_keys(checkpoint_dir: str) -> list: + import glob + + import safetensors + + all_keys = [] + for f in sorted(glob.glob(os.path.join(checkpoint_dir, "model-*.safetensors"))): + with safetensors.safe_open(f, framework="pt") as h: + all_keys.extend(h.keys()) + return all_keys + + +class TestStep3p7Helpers(unittest.TestCase): + """Pure-Python loader/helper tests — no checkpoint and no GPU required.""" + + def test_split_stacked_moe_weights_expands_per_expert_keys(self): + """Stacked routed-expert tensors expand into per-expert ``w1/w2/w3`` keys. + + Source convention: ``...moe.gate_proj.weight`` shape + ``(N, intermediate, hidden)``. Target convention (VANILLA MoE backend): + ``...moe.experts..w1.weight``. Uses a synthetic 2-layer dict (one + MoE, one dense) so the test runs in <1s without the real checkpoint. + """ + from tensorrt_llm._torch.models.modeling_step3p7 import split_stacked_moe_weights + + num_experts = 4 + intermediate = 8 + hidden = 16 + text_config = types.SimpleNamespace( + num_hidden_layers=2, + moe_num_experts=num_experts, + moe_layers_enum=[1], # Only layer 1 is MoE. + ) + + weights = { + # Layer 0 (dense): no MoE stacked tensors. + "model.layers.0.mlp.gate_proj.weight": torch.arange( + intermediate * hidden, dtype=torch.float32 + ).reshape(intermediate, hidden), + # Layer 1 (MoE): stacked routed-expert tensors. + "model.layers.1.moe.gate_proj.weight": torch.arange( + num_experts * intermediate * hidden, dtype=torch.float32 + ).reshape(num_experts, intermediate, hidden), + "model.layers.1.moe.gate_proj.weight_scale_inv": torch.arange( + num_experts * 1 * 2, dtype=torch.float32 + ).reshape(num_experts, 1, 2), + "model.layers.1.moe.up_proj.weight": torch.arange( + num_experts * intermediate * hidden, dtype=torch.float32 + ).reshape(num_experts, intermediate, hidden) + + 1.0, + "model.layers.1.moe.up_proj.weight_scale_inv": torch.arange( + num_experts * 1 * 2, dtype=torch.float32 + ).reshape(num_experts, 1, 2) + + 1.0, + "model.layers.1.moe.down_proj.weight": torch.arange( + num_experts * hidden * intermediate, dtype=torch.float32 + ).reshape(num_experts, hidden, intermediate), + "model.layers.1.moe.down_proj.weight_scale_inv": torch.arange( + num_experts * 2 * 1, dtype=torch.float32 + ).reshape(num_experts, 2, 1), + } + original_gate = weights["model.layers.1.moe.gate_proj.weight"] + original_gate_scale = weights["model.layers.1.moe.gate_proj.weight_scale_inv"] + + layers_split = split_stacked_moe_weights(weights, text_config) + self.assertEqual(layers_split, 1) + + # Stacked keys must be removed. + self.assertNotIn("model.layers.1.moe.gate_proj.weight", weights) + self.assertNotIn("model.layers.1.moe.gate_proj.weight_scale_inv", weights) + self.assertNotIn("model.layers.1.moe.up_proj.weight", weights) + self.assertNotIn("model.layers.1.moe.down_proj.weight", weights) + + # Per-expert keys present with correct slicing. Source ``gate_proj`` → + # ``w1``, ``up_proj`` → ``w3``, ``down_proj`` → ``w2``. + for expert_id in range(num_experts): + key_w1 = f"model.layers.1.moe.experts.{expert_id}.w1.weight" + key_w1_scale = f"model.layers.1.moe.experts.{expert_id}.w1.weight_scale_inv" + key_w3 = f"model.layers.1.moe.experts.{expert_id}.w3.weight" + key_w2 = f"model.layers.1.moe.experts.{expert_id}.w2.weight" + self.assertIn(key_w1, weights) + self.assertIn(key_w1_scale, weights) + self.assertIn(key_w3, weights) + self.assertIn(key_w2, weights) + self.assertTrue(torch.equal(weights[key_w1], original_gate[expert_id])) + self.assertTrue(torch.equal(weights[key_w1_scale], original_gate_scale[expert_id])) + # Per-expert tensor drops the leading expert dim. + self.assertEqual(weights[key_w1].shape, (intermediate, hidden)) + self.assertEqual(weights[key_w2].shape, (hidden, intermediate)) + + # Layer 0 (dense) is untouched. + self.assertIn("model.layers.0.mlp.gate_proj.weight", weights) + + def test_split_stacked_moe_weights_no_op_when_no_stacked_keys(self): + """With no stacked MoE keys present, the splitter is a no-op.""" + from tensorrt_llm._torch.models.modeling_step3p7 import split_stacked_moe_weights + + text_config = types.SimpleNamespace( + num_hidden_layers=1, + moe_num_experts=8, + moe_layers_enum=[0], + ) + weights = {"some.other.key": "not_a_tensor"} + self.assertEqual(split_stacked_moe_weights(weights, text_config), 0) + self.assertEqual(weights, {"some.other.key": "not_a_tensor"}) + + def test_split_stacked_moe_weights_handles_nvfp4_suffixes(self): + """NVFP4 checkpoints carry ``weight_scale`` / ``weight_scale_2`` / + ``input_scale`` alongside the packed ``weight``; the splitter must fan + each suffix out to per-expert keys.""" + from tensorrt_llm._torch.models.modeling_step3p7 import split_stacked_moe_weights + + num_experts = 3 + intermediate = 8 + hidden = 16 + text_config = types.SimpleNamespace( + num_hidden_layers=1, + moe_num_experts=num_experts, + moe_layers_enum=[0], + ) + weights = { + # Packed FP4 weight: (E, I, H/2) uint8. + "model.layers.0.moe.gate_proj.weight": torch.zeros( + (num_experts, intermediate, hidden // 2), dtype=torch.uint8 + ), + # Per-16 block scale: (E, I, H/16) fp8_e4m3fn. + "model.layers.0.moe.gate_proj.weight_scale": torch.ones( + (num_experts, intermediate, hidden // 16), dtype=torch.float8_e4m3fn + ), + # Per-tensor global scale: (E,) float32. + "model.layers.0.moe.gate_proj.weight_scale_2": torch.arange( + num_experts, dtype=torch.float32 + ), + # Per-tensor input scale: (E,) float32. + "model.layers.0.moe.gate_proj.input_scale": torch.arange( + num_experts, dtype=torch.float32 + ) + + 10.0, + } + layers_split = split_stacked_moe_weights(weights, text_config) + self.assertEqual(layers_split, 1) + self.assertNotIn("model.layers.0.moe.gate_proj.weight_scale", weights) + self.assertNotIn("model.layers.0.moe.gate_proj.weight_scale_2", weights) + self.assertNotIn("model.layers.0.moe.gate_proj.input_scale", weights) + for e in range(num_experts): + self.assertEqual( + weights[f"model.layers.0.moe.experts.{e}.w1.weight"].shape, + (intermediate, hidden // 2), + ) + self.assertEqual( + weights[f"model.layers.0.moe.experts.{e}.w1.weight_scale"].shape, + (intermediate, hidden // 16), + ) + self.assertEqual( + weights[f"model.layers.0.moe.experts.{e}.w1.weight_scale_2"].item(), float(e) + ) + self.assertEqual( + weights[f"model.layers.0.moe.experts.{e}.w1.input_scale"].item(), float(e) + 10.0 + ) + + def test_rewrite_mtp_weights_uses_mtp_block_layout(self): + """Checkpoint MTP keys map to TRT-LLM's ``mtp_block`` layout.""" + from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict + from tensorrt_llm._torch.models.modeling_step3p7 import rewrite_mtp_weights_for_step3p7 + + # The rewriter only relies on ``num_hidden_layers`` (45 decoder layers, + # so MTP layers start at index 45) and ``num_nextn_predict_layers`` + # (3 MTP layers); a lightweight stub avoids loading the real checkpoint. + text_config = types.SimpleNamespace(num_hidden_layers=45, num_nextn_predict_layers=3) + weights = ConsumableWeightsDict( + { + "model.layers.45.enorm.weight": torch.tensor([1.0]), + "model.layers.45.eh_proj.weight": torch.tensor([2.0]), + "model.layers.45.self_attn.q_proj.weight": torch.tensor([3.0]), + "model.layers.45.mlp.gate_proj.weight": torch.tensor([4.0]), + "model.layers.45.transformer.shared_head.norm.weight": torch.tensor([5.0]), + "model.layers.45.transformer.shared_head.output.weight": torch.tensor([6.0]), + "model.layers.44.self_attn.q_proj.weight": torch.tensor([7.0]), + } + ) + + rewritten = rewrite_mtp_weights_for_step3p7(weights, text_config) + + self.assertEqual(rewritten, 4) + self.assertIn("model.layers.45.enorm.weight", weights) + self.assertIn("model.layers.45.eh_proj.weight", weights) + self.assertIn("model.layers.45.mtp_block.self_attn.q_proj.weight", weights) + self.assertIn("model.layers.45.mtp_block.mlp.gate_proj.weight", weights) + self.assertIn("model.layers.45.shared_head.norm.weight", weights) + self.assertIn("model.layers.45.shared_head.output.weight", weights) + self.assertIn("model.layers.44.self_attn.q_proj.weight", weights) + self.assertNotIn("model.layers.45.self_attn.q_proj.weight", weights) + self.assertNotIn("model.layers.45.transformer.shared_head.output.weight", weights) + + def test_rewrite_language_model_keys_flattens_multimodal_namespace(self): + """The NVFP4 multimodal checkpoint stores text decoder keys under + ``model.language_model.*``; the loader normalizes them to ``model.*`` so + the module tree (built around ``model.layers.*``) is addressable. Vision + keys move out from under ``model.`` to match the ignored prefix list.""" + from tensorrt_llm._torch.models.modeling_step3p7 import rewrite_language_model_keys + + weights = { + "lm_head.weight": "lm", + "model.language_model.embed_tokens.weight": "embed", + "model.language_model.layers.0.input_layernorm.weight": "in_ln", + "model.language_model.norm.weight": "final_ln", + "model.vision_model.conv1.weight": "vision", + "model.vit_large_projector.weight": "projector", + } + n = rewrite_language_model_keys(weights) + # 5 of 6 keys get renamed; lm_head is untouched. + self.assertEqual(n, 5) + self.assertIn("lm_head.weight", weights) + self.assertIn("model.embed_tokens.weight", weights) + self.assertIn("model.layers.0.input_layernorm.weight", weights) + self.assertIn("model.norm.weight", weights) + self.assertIn("vision_model.conv1.weight", weights) + self.assertIn("vit_large_projector.weight", weights) + self.assertNotIn("model.language_model.embed_tokens.weight", weights) + self.assertNotIn("model.vision_model.conv1.weight", weights) + + def test_strip_language_model_prefix_preserves_regex_and_non_matching(self): + """The quant-config exclude-modules normalizer must replace + ``model.language_model.`` with ``model.``, leave ``re:`` prefixed + entries untouched, and leave entries without the segment untouched.""" + from tensorrt_llm._torch.models.modeling_step3p7 import ( + strip_language_model_prefix_from_exclude_modules, + ) + + src = [ + "lm_head", + "model.language_model.layers.0*", + "model.language_model.layers.3.moe.gate", + "re:^model\\.something\\..*", + "model.vision_model*", + ] + out = strip_language_model_prefix_from_exclude_modules(src) + self.assertEqual( + out, + [ + "lm_head", + "model.layers.0*", + "model.layers.3.moe.gate", + "re:^model\\.something\\..*", + "model.vision_model*", + ], + ) + self.assertIsNone(strip_language_model_prefix_from_exclude_modules(None)) + + def test_nvfp4_dequant_batched_round_trips_constant_values(self): + """Sanity check the NVFP4 dequant helper. Packing a constant e2m1 index + of 4 (= 2.0) and scaling with block_scale=1.0 and global_scale=0.5 must + produce 1.0 everywhere.""" + from tensorrt_llm._torch.models.modeling_step3p7 import _nvfp4_dequant_batched + + # Each byte holds (high<<4) | low. Encoding index 4 in both nibbles: + # (4 << 4) | 4 = 68. + weight = torch.full((2, 8), 68, dtype=torch.uint8) + block_scale = torch.ones((2, 1), dtype=torch.float8_e4m3fn) + global_scale = torch.tensor(0.5, dtype=torch.float32) + out = _nvfp4_dequant_batched(weight, block_scale, global_scale) + self.assertEqual(out.dtype, torch.bfloat16) + # 2.0 (e2m1) * 1.0 (block) * 0.5 (global) = 1.0 + self.assertTrue(torch.all(out == 1.0)) + + # Encode a negative index (12 = -2.0): byte = (12 << 4) | 4 → high=-2, + # low=2. Use K=16 so the per-16 block-scale shape works. + weight2 = torch.full((1, 8), (12 << 4) | 4, dtype=torch.uint8) + out2 = _nvfp4_dequant_batched( + weight2, + torch.ones((1, 1), dtype=torch.float8_e4m3fn), + torch.tensor(1.0, dtype=torch.float32), + ) + # Even positions get 2.0 (low nibble), odd positions get -2.0 (high nibble). + self.assertEqual(out2[0, 0].item(), 2.0) + self.assertEqual(out2[0, 1].item(), -2.0) + + # 3D batched form: (E=2, M=2, K_half=8) — exercises the per-expert + # global-scale broadcast path used by the Python-clamp loader. + w3 = torch.full((2, 2, 8), 68, dtype=torch.uint8) + s1 = torch.ones((2, 2, 1), dtype=torch.float8_e4m3fn) + s2 = torch.tensor([0.5, 0.25], dtype=torch.float32) + out3 = _nvfp4_dequant_batched(w3, s1, s2) + self.assertTrue(torch.all(out3[0] == 1.0)) # 2.0 * 1.0 * 0.5 = 1.0 + self.assertTrue(torch.all(out3[1] == 0.5)) # 2.0 * 1.0 * 0.25 = 0.5 + + def test_mtp_head_normalizes_before_output_projection(self): + """Step3p7 MTP applies shared-head norm only when producing draft logits.""" + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_step3p7 import Step3p7MTPHead + from tensorrt_llm._torch.modules import rms_norm as rms_norm_module + + flashinfer_available = rms_norm_module.IS_FLASHINFER_AVAILABLE + rms_norm_module.IS_FLASHINFER_AVAILABLE = False + try: + text_config = PretrainedConfig() + text_config.hidden_size = 2 + text_config.rms_norm_eps = 0.0 + text_config.torch_dtype = torch.float32 + text_config.vocab_size = 2 + top_config = PretrainedConfig() + top_config.text_config = text_config + head = Step3p7MTPHead(ModelConfig(pretrained_config=top_config)) + + class CaptureOutput(torch.nn.Module): + gather_output = True + + def __init__(self): + super().__init__() + self.seen = None + + def forward(self, hidden_states, **kwargs): + del kwargs + self.seen = hidden_states.detach().clone() + return hidden_states + + output = CaptureOutput() + head.output = output + hidden_states = torch.tensor([[3.0, 4.0], [6.0, 8.0]], dtype=torch.float32) + + logits = head(hidden_states, lm_head=None, attn_metadata=None) + + expected = hidden_states[-1:] / hidden_states[-1:].pow(2).mean(-1, keepdim=True).sqrt() + self.assertTrue(torch.allclose(output.seen, expected)) + self.assertTrue(torch.allclose(logits, expected)) + finally: + rms_norm_module.IS_FLASHINFER_AVAILABLE = flashinfer_available + + +class TestStep3p7AutoModelRegistration(unittest.TestCase): + """Verify the Step3p7 architectures resolve to the expected model classes.""" + + def test_auto_model_registered(self): + from tensorrt_llm._torch.models.modeling_step3p7 import Step3p7ForCausalLM + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3p7VLForConditionalGeneration + from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING + + # Step3p7ForConditionalGeneration is the multimodal entry point; the + # text-only causal LM remains reachable via "Step3p5ForCausalLM". + self.assertIs( + MODEL_CLASS_MAPPING.get("Step3p7ForConditionalGeneration"), + Step3p7VLForConditionalGeneration, + ) + self.assertIs(MODEL_CLASS_MAPPING.get("Step3p5ForCausalLM"), Step3p7ForCausalLM) + + +class TestStep3p7Checkpoint(unittest.TestCase): + """Config / weight-accounting tests against the real Step-3.7-Flash checkpoints. + + The FP8 block-scale, NVFP4, and BF16 reference checkpoints share the same + per-layer geometry; only the routed-expert dtype/layout differs. The FP8 and + BF16 checkpoints store text keys under ``model.layers.*`` with 3 plain-path + MTP layers; the NVFP4 modelopt export stores text keys under + ``model.language_model.layers.*`` and carries no MTP layers. The checkpoints + are expected under ``LLM_MODELS_ROOT`` (as in CI); a missing checkpoint + surfaces as a test failure rather than a silent skip. + """ + + def _check_text_config(self, text_cfg: dict): + self.assertEqual(text_cfg["model_type"], "step3p5") + self.assertEqual(text_cfg["num_hidden_layers"], 45) + self.assertEqual(text_cfg["hidden_size"], 4096) + self.assertEqual(text_cfg["vocab_size"], 128896) + self.assertEqual(text_cfg["num_attention_heads"], 64) + self.assertEqual(text_cfg["num_attention_groups"], 8) + self.assertEqual(text_cfg["head_dim"], 128) + self.assertEqual(text_cfg["moe_num_experts"], 288) + self.assertEqual(text_cfg["moe_top_k"], 8) + self.assertEqual(text_cfg["moe_router_scaling_factor"], 3.0) + self.assertIs(text_cfg["use_head_wise_attn_gate"], True) + self.assertIs(text_cfg["use_moe_router_bias"], True) + self.assertIs(text_cfg["need_fp32_gate"], True) + + @parameterized.expand( + [ + ("fp8", STEP3P7_FP8_DIR), + ("nvfp4", STEP3P7_NVFP4_DIR), + ("bf16", STEP3P7_BF16_DIR), + ] + ) + def test_config_and_weight_accounting(self, name, checkpoint_dir): + """Recognize Step3p7 + account for every safetensors key. + + - architectures == ["Step3p7ForConditionalGeneration"], top + model_type == "step3p7", text model_type == "step3p5" + - 45 text decoder layers with the documented full/sliding pattern + - FP8 (fp8 block-scale) and NVFP4 (modelopt) checkpoints carry a + quantization_config block; the BF16 reference does not + - All consumed text-path keys can be enumerated; only vision, plain-path + MTP (layers 45..47, FP8/BF16 only), and the vision projector are + ignored. The NVFP4 export nests text keys under + ``model.language_model.*`` and vision keys under ``model.vision_model.*``. + """ + import re + + is_fp8 = name == "fp8" + is_nvfp4 = name == "nvfp4" + config_dict = _load_config(checkpoint_dir) + safetensors_keys = _load_safetensors_keys(checkpoint_dir) + + # 1. Top-level config sanity. + self.assertEqual(config_dict["architectures"], ["Step3p7ForConditionalGeneration"]) + self.assertEqual(config_dict["model_type"], "step3p7") + text_cfg = config_dict["text_config"] + self._check_text_config(text_cfg) + + # 2. Layer inventory: 45 decoder layers, full at idx 0,4,8,...,44 and + # sliding elsewhere. The raw layer_types array can be longer (48 + # entries) because it also covers the 3 MTP layers (45..47); the + # NVFP4 export has no MTP layers so it carries exactly 45 entries. + layer_types = text_cfg["layer_types"] + self.assertGreaterEqual(len(layer_types), 45) + for idx in range(45): + lt = layer_types[idx] + if idx % 4 == 0: + self.assertEqual(lt, "full_attention", f"layer {idx}") + else: + self.assertEqual(lt, "sliding_attention", f"layer {idx}") + + # 3. Quant config: FP8 and NVFP4 carry it, BF16 does not. The NVFP4 + # export uses the modelopt schema (quant_algo + ``ignore`` list with + # ``model.language_model.*`` patterns) rather than the fp8 schema + # (weight_block_size + ``modules_to_not_convert``). + if is_fp8: + quant_cfg = config_dict["quantization_config"] + self.assertEqual(quant_cfg["quant_method"], "fp8") + self.assertEqual(quant_cfg["weight_block_size"], [128, 128]) + not_convert = set(quant_cfg["modules_to_not_convert"]) + for layer_idx in range(45): + if layer_idx < 3: + # Dense MLP layers: gate/up/down all bf16. + for sub in ( + "self_attn.q_proj", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.down_proj", + ): + key = f"model.layers.{layer_idx}.{sub}" + self.assertIn(key, not_convert, f"dense layer {layer_idx} {sub}") + else: + # MoE layers: routed gate/up/down are FP8; shared expert is bf16. + for sub in ( + "self_attn.q_proj", + "moe.gate", + "share_expert.gate_proj", + "share_expert.up_proj", + "share_expert.down_proj", + ): + key = f"model.layers.{layer_idx}.{sub}" + self.assertIn(key, not_convert, f"MoE layer {layer_idx} {sub}") + elif is_nvfp4: + quant_cfg = config_dict["quantization_config"] + self.assertEqual(quant_cfg["quant_method"], "modelopt") + self.assertEqual(quant_cfg["quant_algo"], "NVFP4") + # Routed-expert matmuls quantize to NVFP4; everything else (lm_head, + # router gate, shared expert, attention) is excluded via the + # ``ignore`` list keyed off the ``model.language_model.*`` namespace. + ignore = set(quant_cfg["ignore"]) + self.assertIn("lm_head", ignore) + self.assertTrue( + any(e.startswith("model.language_model.layers.") for e in ignore), + "expected language_model exclude entries in NVFP4 ignore list", + ) + else: + self.assertNotIn("quantization_config", config_dict) + + # 4. Weight accounting on the actual safetensors keys. The NVFP4 export + # nests the text decoder under ``model.language_model.*`` and the + # vision tower under ``model.vision_model.*``; FP8/BF16 use bare + # ``model.layers.*`` / ``vision_model.*`` prefixes. + if is_nvfp4: + vision_prefixes = ("model.vision_model.", "model.vit_large_projector") + embed_norm_keys = ( + "model.language_model.embed_tokens.weight", + "model.language_model.norm.weight", + "lm_head.weight", + ) + text_layer_re = re.compile(r"^model\.language_model\.layers\.(\d+)\.") + else: + vision_prefixes = ("vision_model.", "vit_large_projector") + embed_norm_keys = ( + "model.embed_tokens.weight", + "model.norm.weight", + "lm_head.weight", + ) + text_layer_re = re.compile(r"^model\.layers\.(\d+)\.") + + consumed_text_keys = [] + ignored_vision_keys = [] + ignored_mtp_keys = [] + for key in safetensors_keys: + m = text_layer_re.match(key) + if m: + if int(m.group(1)) >= 45: + ignored_mtp_keys.append(key) + else: + consumed_text_keys.append(key) + elif key.startswith(vision_prefixes): + ignored_vision_keys.append(key) + elif key in embed_norm_keys: + consumed_text_keys.append(key) + else: + self.fail(f"Unaccounted-for safetensors key: {key}") + + self.assertGreater(len(consumed_text_keys), 0, "no consumed text keys found") + self.assertGreater(len(ignored_vision_keys), 0, "no ignored vision keys found") + # FP8/BF16 ship 3 plain-path MTP layers (45..47); the NVFP4 export omits + # them entirely. + if is_nvfp4: + self.assertEqual(len(ignored_mtp_keys), 0, "NVFP4 export should carry no MTP keys") + else: + self.assertGreater(len(ignored_mtp_keys), 0, "no ignored MTP keys (expected 45..47)") + + @parameterized.expand( + [ + ("fp8", STEP3P7_FP8_DIR), + ("nvfp4", STEP3P7_NVFP4_DIR), + ("bf16", STEP3P7_BF16_DIR), + ] + ) + def test_model_config_resolves_registered_architecture(self, name, checkpoint_dir): + """``ModelConfig.from_pretrained`` resolves Step3p7 and detects the quant algo.""" + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.quantization import QuantAlgo + + model_config = ModelConfig.from_pretrained(checkpoint_dir, trust_remote_code=True) + pc = model_config.pretrained_config + self.assertEqual(pc.architectures, ["Step3p7ForConditionalGeneration"]) + self.assertEqual(pc.num_hidden_layers, 45) + self.assertEqual(pc.vocab_size, 128896) + self.assertEqual(pc.hidden_size, 4096) + if name == "fp8": + self.assertEqual(model_config.quant_config.quant_algo, QuantAlgo.FP8_BLOCK_SCALES) + elif name == "nvfp4": + self.assertEqual(model_config.quant_config.quant_algo, QuantAlgo.NVFP4) + # NVFP4 modelopt export also carries an FP8 KV-cache scheme. + self.assertEqual(model_config.quant_config.kv_cache_quant_algo, QuantAlgo.FP8) + else: + self.assertTrue( + model_config.quant_config is None or model_config.quant_config.quant_algo is None + ) + + @parameterized.expand( + [ + ("fp8", STEP3P7_FP8_DIR), + ("nvfp4", STEP3P7_NVFP4_DIR), + ("bf16", STEP3P7_BF16_DIR), + ] + ) + def test_per_layer_attention_geometry_helpers(self, name, checkpoint_dir): + """Per-layer head / RoPE / SwiGLU helpers match config arithmetic. + + Runs against all three checkpoints because the per-layer attention + geometry is identical between them: only the routed-expert dtype/layout + differs. The layer-45 assertions exercise the out-of-range fallback path + (the NVFP4 export carries only 45 ``layer_types`` entries; FP8/BF16 + include the 3 MTP layers), which resolves to the same sliding-attention + geometry either way. + """ + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_step3p7 import ( + _is_moe_layer, + _layer_attention_type, + _layer_kv_heads, + _layer_partial_rotary, + _layer_query_heads, + _layer_rope_theta, + _layer_swiglu_limit, + _layer_uses_rope_scaling, + ) + + model_config = ModelConfig.from_pretrained(checkpoint_dir, trust_remote_code=True) + text_config = model_config.pretrained_config.text_config + + # Full vs sliding head counts. + self.assertEqual(_layer_query_heads(text_config, 0), 64) + self.assertEqual(_layer_kv_heads(text_config, 0), 8) + self.assertEqual(_layer_query_heads(text_config, 1), 96) + self.assertEqual(_layer_kv_heads(text_config, 1), 8) + self.assertEqual(_layer_attention_type(text_config, 0), "full_attention") + self.assertEqual(_layer_attention_type(text_config, 1), "sliding_attention") + self.assertEqual(_layer_attention_type(text_config, 45), "sliding_attention") + self.assertEqual(_layer_query_heads(text_config, 45), 96) + self.assertEqual(_layer_kv_heads(text_config, 45), 8) + + # RoPE per-layer. + self.assertEqual(_layer_rope_theta(text_config, 0), 5_000_000.0) + self.assertEqual(_layer_rope_theta(text_config, 1), 10_000.0) + self.assertEqual(_layer_rope_theta(text_config, 45), 10_000.0) + self.assertEqual(_layer_partial_rotary(text_config, 0), 0.5) + self.assertEqual(_layer_partial_rotary(text_config, 1), 1.0) + self.assertEqual(_layer_partial_rotary(text_config, 45), 1.0) + self.assertIs(_layer_uses_rope_scaling(text_config, 0), True) + self.assertIs(_layer_uses_rope_scaling(text_config, 1), False) + + # MoE vs dense. + self.assertIs(_is_moe_layer(text_config, 0), False) + self.assertIs(_is_moe_layer(text_config, 1), False) + self.assertIs(_is_moe_layer(text_config, 2), False) + self.assertIs(_is_moe_layer(text_config, 3), True) + self.assertIs(_is_moe_layer(text_config, 44), True) + + # SwiGLU clamp limits: nonzero only on layers 43 and 44 (per config). + self.assertIsNone(_layer_swiglu_limit(text_config, 0)) + self.assertEqual(_layer_swiglu_limit(text_config, 43), 7.0) + self.assertEqual(_layer_swiglu_limit(text_config, 44), 7.0) + self.assertIsNone(_layer_swiglu_limit(text_config, 45)) + self.assertEqual(_layer_swiglu_limit(text_config, 43, shared=True), 16.0) + self.assertEqual(_layer_swiglu_limit(text_config, 44, shared=True), 16.0) + + def test_mtp_spec_config_defaults_to_checkpoint_layer_count(self): + """Default ``MTPDecodingConfig`` loads all Step3p7 MTP layers.""" + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_step3p7 import _prepare_step3p7_mtp_spec_config + from tensorrt_llm._torch.speculative import update_spec_config_from_model_config + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + + spec_config = MTPDecodingConfig() + # Unset → None sentinel; resolved to the checkpoint layer count below. + self.assertIsNone(spec_config.max_draft_len) + self.assertNotIn("max_draft_len", spec_config.model_fields_set) + model_config = ModelConfig.from_pretrained( + STEP3P7_FP8_DIR, trust_remote_code=True, spec_config=spec_config + ) + update_spec_config_from_model_config(spec_config, model_config.pretrained_config) + _prepare_step3p7_mtp_spec_config(model_config) + + self.assertEqual(spec_config.num_nextn_predict_layers, 3) + self.assertEqual(spec_config.max_draft_len, 3) + self.assertEqual(spec_config.max_total_draft_tokens, 3) + self.assertTrue(spec_config.spec_dec_mode.is_mtp_vanilla()) + + explicit_spec_config = MTPDecodingConfig(max_draft_len=1) + explicit_model_config = ModelConfig.from_pretrained( + STEP3P7_FP8_DIR, trust_remote_code=True, spec_config=explicit_spec_config + ) + update_spec_config_from_model_config( + explicit_spec_config, explicit_model_config.pretrained_config + ) + _prepare_step3p7_mtp_spec_config(explicit_model_config) + + self.assertEqual(explicit_spec_config.num_nextn_predict_layers, 3) + self.assertEqual(explicit_spec_config.max_draft_len, 1) + self.assertEqual(explicit_spec_config.max_total_draft_tokens, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/modeling/test_modeling_step3p7vl.py b/tests/unittest/_torch/modeling/test_modeling_step3p7vl.py new file mode 100644 index 000000000000..ba1edc539dd4 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_modeling_step3p7vl.py @@ -0,0 +1,529 @@ +# 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. +"""Tests for the Step3p7 multimodal bring-up (``modeling_step3p7vl``). + +The text decoder + MTP wiring is covered by ``test_modeling_step3p7.py``; this +module focuses on the Perception-Encoder vision tower and the VLM registration. + +TestStep3p7VisionTower — vision tower component tests (no checkpoint, no GPU). +A shrunken synthetic ``vision_config`` (see ``_make_tiny_vision_config``) keeps +the real per-layer geometry (Conv2d patch embed -> pre-LN transformer blocks +with 2D RoPE + LayerScale -> two stride-2 Conv2d downsamplers -> linear +projector) while letting the tower build and run on CPU in well under a second: + - 2D RoPE helpers (zero-frequency identity at grid position 0, freq-cache + shape, dynamic sub-grid selection) + - LayerScale / fused-QKV attention / MLP / pre-LN block component contracts + - ``Step3p7VisionEncoder`` forward output shape, including the abs-posemb + bilinear interpolation path for off-grid image sizes + - ``Step3p7VisionTower`` projector geometry, ``vision_model.*`` / + ``vit_large_projector.*`` weight-prefix splitting, and the per-request + ``[patches... | full image]`` embedding flattening contract + +TestStep3p7VLRegistration — registry verification: + - ``Step3p7ForConditionalGeneration`` resolves to the VLM entry point + - ``Step3p7VisionTower`` is registered as the architecture's vision encoder + - the multimodal input processor is registered for ``model_type == step3p7`` + +TestStep3p7VLCheckpoint — checkpoint-backed ``vision_config`` geometry tests +(requires the Step-3.7-Flash checkpoints under ``LLM_MODELS_ROOT``). The FP8 +block-scale, NVFP4, and BF16 reference checkpoints all ship the same +PerceptionEncoder vision config; only the routed-expert dtype/layout of the +text decoder differs. +""" + +import json +import os +import unittest + +import torch +from parameterized import parameterized +from PIL import Image +from transformers import PretrainedConfig +from utils.llm_data import llm_models_root + +# The multimodal checkpoint is the same on disk as the text one; the FP8 +# block-scale, NVFP4, and BF16 reference checkpoints all ship the same +# PerceptionEncoder vision tower (only the text decoder's routed-expert +# dtype/layout differs). See test_modeling_step3p7.py for the text-path tests. +STEP3P7_FP8_DIR = str(os.path.join(llm_models_root(), "Step-3.7-Flash-FP8")) +STEP3P7_NVFP4_DIR = str(os.path.join(llm_models_root(), "Step-3.7-Flash-NVFP4")) +STEP3P7_BF16_DIR = str(os.path.join(llm_models_root(), "Step-3.7-Flash")) + + +def _load_config(checkpoint_dir: str) -> dict: + with open(os.path.join(checkpoint_dir, "config.json")) as f: + return json.load(f) + + +def _make_tiny_vision_config( + width: int = 64, + heads: int = 4, + layers: int = 2, + patch_size: int = 8, + image_size: int = 64, + ls_init_value: float = 0.1, + use_cls_token: bool = False, + use_ln_post: bool = False, + hidden_act: str = "quick_gelu", +) -> PretrainedConfig: + """Tiny ``vision_config`` mirroring the Step-3.7-Flash PerceptionEncoder. + + Field names match the real checkpoint (``width`` / ``heads`` / ``layers`` + rather than ``hidden_size`` / ``num_heads`` / ``num_hidden_layers``) so the + encoder reads them exactly as it would from ``config.json``. Sizes are + shrunk so the tower builds and runs on CPU in well under a second; float32 + keeps the identity invariants below numerically exact. + """ + cfg = PretrainedConfig() + cfg.model_type = "perception_encoder" + cfg.width = width + cfg.heads = heads + cfg.layers = layers + cfg.patch_size = patch_size + cfg.image_size = image_size + cfg.hidden_act = hidden_act + cfg.ls_init_value = ls_init_value + cfg.use_cls_token = use_cls_token + cfg.use_ln_post = use_ln_post + cfg.torch_dtype = torch.float32 + return cfg + + +def _make_tiny_vision_model_config(text_hidden_size: int = 32, **vision_kwargs): + """Wrap a tiny vision config in a ``ModelConfig`` the vision tower accepts. + + ``Step3p7VisionTower`` reads ``vision_config``, ``text_config.hidden_size``, + ``torch_dtype``, ``image_token_id`` and ``projector_bias`` off the top-level + pretrained config; supply just those. + """ + from tensorrt_llm._torch.model_config import ModelConfig + + vision_cfg = _make_tiny_vision_config(**vision_kwargs) + top = PretrainedConfig() + top.torch_dtype = torch.float32 + top.vision_config = vision_cfg + text_cfg = PretrainedConfig() + text_cfg.hidden_size = text_hidden_size + top.text_config = text_cfg + top.image_token_id = 128001 + top.projector_bias = False + return ModelConfig(pretrained_config=top) + + +class TestStep3p7VisionTower(unittest.TestCase): + """Perception-Encoder vision tower tests — no checkpoint and no GPU.""" + + @staticmethod + def _downsampled_tokens(grid: int) -> int: + """Token count after the two trailing stride-2 / kernel-3 / pad-1 convs. + + Each downsampler maps spatial ``s -> floor((s - 1) / 2) + 1``; applied + twice this equals ``grid // 4`` for the grid sizes used here, matching + the encoder's documented ``(Gh//4) * (Gw//4)`` output. + """ + after1 = (grid - 1) // 2 + 1 + after2 = (after1 - 1) // 2 + 1 + return after2 * after2 + + # ----- 2D RoPE ------------------------------------------------------- + + def test_rope2d_freqs_cache_shape(self): + """Cached frequencies are ``(1, 1, Gh*Gw, head_dim)``.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3VisionRope2D + + head_dim, gh, gw = 16, 4, 4 + rope = Step3VisionRope2D(dim=head_dim, max_grid_height=gh, max_grid_width=gw) + self.assertEqual(tuple(rope.freqs_cache.shape), (1, 1, gh * gw, head_dim)) + + def test_rope2d_position_zero_is_identity(self): + """At grid position (0, 0) the 2D frequencies are zero → cos=1, sin=0. + + The first sequence position must therefore pass through unchanged while + a later position is actually rotated. + """ + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3VisionRope2D + + head_dim, gh, gw = 16, 4, 4 + rope = Step3VisionRope2D(dim=head_dim, max_grid_height=gh, max_grid_width=gw) + q = torch.randn(1, 2, gh * gw, head_dim) + k = torch.randn(1, 2, gh * gw, head_dim) + q_out, k_out = rope(q, k, grid_hw=(gh, gw)) + self.assertEqual(q_out.shape, q.shape) + self.assertTrue(torch.allclose(q_out[..., 0, :], q[..., 0, :], atol=1e-6)) + self.assertTrue(torch.allclose(k_out[..., 0, :], k[..., 0, :], atol=1e-6)) + self.assertFalse(torch.allclose(q_out[..., -1, :], q[..., -1, :], atol=1e-4)) + + def test_rope2d_dynamic_subgrid_selects_positions(self): + """A grid smaller than the cached max grid hits the index-select path.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3VisionRope2D + + head_dim = 16 + rope = Step3VisionRope2D(dim=head_dim, max_grid_height=8, max_grid_width=8) + gh, gw = 4, 4 + q = torch.randn(1, 2, gh * gw, head_dim) + k = torch.randn(1, 2, gh * gw, head_dim) + q_out, k_out = rope(q, k, grid_hw=(gh, gw)) + self.assertEqual(q_out.shape, (1, 2, gh * gw, head_dim)) + self.assertEqual(k_out.shape, (1, 2, gh * gw, head_dim)) + + # ----- components ---------------------------------------------------- + + def test_layer_scale_scales_by_gamma(self): + """LayerScale multiplies its input element-wise by the per-channel gamma.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3VisionLayerScale + + ls = Step3VisionLayerScale(dim=4, init_value=2.0) + x = torch.ones(1, 3, 4) + self.assertTrue(torch.allclose(ls(x), x * 2.0)) + + def test_vision_mlp_names_and_shape(self): + """The FFN keeps the HF ``c_fc`` / ``c_proj`` parameter names.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3VisionMLP + + mlp = Step3VisionMLP(hidden_size=8, intermediate_size=16, hidden_act="quick_gelu") + self.assertTrue(hasattr(mlp, "c_fc")) + self.assertTrue(hasattr(mlp, "c_proj")) + out = mlp(torch.randn(2, 5, 8)) + self.assertEqual(out.shape, (2, 5, 8)) + + def test_vision_attention_fused_qkv_layout_and_shape(self): + """HF fused-QKV layout (``in_proj_weight`` / ``in_proj_bias``) and shape.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3VisionAttention + + hidden, heads, gh, gw = 64, 4, 4, 4 + attn = Step3VisionAttention( + hidden_size=hidden, + num_heads=heads, + max_grid_height=gh, + max_grid_width=gw, + use_cls_token=False, + use_rope2d=True, + ) + self.assertEqual(tuple(attn.in_proj_weight.shape), (3 * hidden, hidden)) + self.assertEqual(tuple(attn.in_proj_bias.shape), (3 * hidden,)) + out = attn(torch.randn(2, gh * gw, hidden), grid_hw=(gh, gw)) + self.assertEqual(out.shape, (2, gh * gw, hidden)) + + def test_vision_attention_rejects_indivisible_heads(self): + """``hidden_size`` not divisible by ``num_heads`` is rejected up front.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3VisionAttention + + with self.assertRaises(ValueError): + Step3VisionAttention( + hidden_size=65, + num_heads=4, + max_grid_height=4, + max_grid_width=4, + use_cls_token=False, + use_rope2d=False, + ) + + def test_vision_block_zero_layerscale_is_identity(self): + """With ``ls_init_value=0`` both residual branches are scaled to zero, so + a pre-LN block reduces to the identity.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3VisionBlock + + hidden, heads, gh, gw = 64, 4, 4, 4 + block = Step3VisionBlock( + hidden_size=hidden, + num_heads=heads, + mlp_ratio=2.0, + hidden_act="quick_gelu", + layer_norm_eps=1e-5, + ls_init_value=0.0, + max_grid_height=gh, + max_grid_width=gw, + use_cls_token=False, + use_rope2d=True, + rope_theta=10000.0, + rope_theta_rescale_factor=1.0, + ) + x = torch.randn(1, gh * gw, hidden) + out = block(x, grid_hw=(gh, gw)) + self.assertTrue(torch.allclose(out, x, atol=1e-6)) + + # ----- encoder ------------------------------------------------------- + + def test_vision_encoder_forward_output_shape(self): + """Encoder returns ``(B, (Gh//4)*(Gw//4), 4*width)`` post-downsample.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3p7VisionEncoder + + width, patch, image = 64, 8, 64 + enc = Step3p7VisionEncoder(_make_tiny_vision_config(width=width), dtype=torch.float32) + with torch.inference_mode(): + feats = enc(torch.randn(1, 3, image, image)) + grid = image // patch + self.assertEqual(feats.shape, (1, self._downsampled_tokens(grid), 4 * width)) + + def test_vision_encoder_smaller_image_interpolates_posemb(self): + """An off-grid (smaller) image triggers bilinear abs-posemb interpolation + and the RoPE sub-grid path, still yielding a well-formed feature map.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3p7VisionEncoder + + width, patch, image = 64, 8, 128 # base grid 16 + enc = Step3p7VisionEncoder( + _make_tiny_vision_config(width=width, patch_size=patch, image_size=image), + dtype=torch.float32, + ) + smaller = 64 # grid 8 < base grid 16 + with torch.inference_mode(): + feats = enc(torch.randn(1, 3, smaller, smaller)) + grid = smaller // patch + self.assertEqual(feats.shape, (1, self._downsampled_tokens(grid), 4 * width)) + + # ----- tower (encoder + projector) ----------------------------------- + + def test_vision_tower_projector_geometry(self): + """Projector maps ``4*width -> text hidden_size``; tower honours dtype.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3p7VisionTower + + tower = Step3p7VisionTower(_make_tiny_vision_model_config(text_hidden_size=32, width=64)) + self.assertEqual(tower.vit_large_projector.in_features, 4 * 64) + self.assertEqual(tower.vit_large_projector.out_features, 32) + self.assertIsNone(tower.vit_large_projector.bias) # projector_bias=False + self.assertEqual(tower.dtype, torch.float32) + + def test_vision_tower_encode_projects_to_text_hidden(self): + """``_encode`` runs the encoder + projector to the text hidden size.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3p7VisionTower + + tower = Step3p7VisionTower( + _make_tiny_vision_model_config( + text_hidden_size=32, width=64, patch_size=8, image_size=64 + ) + ) + with torch.inference_mode(): + out = tower._encode(torch.randn(1, 3, 64, 64)) + self.assertEqual(out.shape, (1, self._downsampled_tokens(64 // 8), 32)) + + def test_vision_tower_forward_flattens_patches_then_image(self): + """``forward`` lays each request out as ``[patches... | full image]``. + + Matches the input processor's placeholder expansion: per image, the + per-patch feature blocks come first (in order), then the full-image + block, all flattened to ``(num_tokens, text_hidden)``. + """ + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3p7VisionTower + from tensorrt_llm.inputs.multimodal import MultimodalParams + + tower = Step3p7VisionTower( + _make_tiny_vision_model_config( + text_hidden_size=32, width=64, patch_size=8, image_size=64 + ) + ) + tokens_per_tile = self._downsampled_tokens(64 // 8) + + # One full image, no patches → just the full-image block. + mm_image_only = MultimodalParams( + multimodal_data={"image": {"pixel_values": torch.randn(1, 3, 64, 64)}} + ) + out = tower.forward([mm_image_only]) + self.assertEqual(len(out), 1) + self.assertEqual(out[0].shape, (tokens_per_tile, 32)) + + # One image with 2 patches → 2 patch blocks then the full-image block. + mm_with_patches = MultimodalParams( + multimodal_data={ + "image": { + "pixel_values": torch.randn(1, 3, 64, 64), + "patch_pixel_values": torch.randn(2, 3, 64, 64), + "num_patches": [2], + } + } + ) + out = tower.forward([mm_with_patches]) + self.assertEqual(len(out), 1) + self.assertEqual(out[0].shape, (3 * tokens_per_tile, 32)) + + # No image payload → no embeddings produced. + self.assertEqual(tower.forward([MultimodalParams(multimodal_data={})]), []) + + def test_vision_tower_load_weights_splits_prefixes(self): + """``load_weights`` routes ``vision_model.*`` to the encoder and + ``vit_large_projector.*`` to the projector, ignoring unrelated keys.""" + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3p7VisionTower + + tower = Step3p7VisionTower(_make_tiny_vision_model_config(text_hidden_size=32, width=64)) + # ``load_weights`` loads each subtree with ``strict=True``, so supply a + # complete vision/projector state built from the modules' own keys, then + # override two routed tensors with sentinels to assert correct routing. + weights = { + f"vision_model.{k}": v.clone() for k, v in tower.vision_model.state_dict().items() + } + weights.update( + { + f"vit_large_projector.{k}": v.clone() + for k, v in tower.vit_large_projector.state_dict().items() + } + ) + proj_w = torch.ones_like(tower.vit_large_projector.weight) + conv_w = torch.ones_like(tower.vision_model.conv1.weight) + weights["vit_large_projector.weight"] = proj_w + weights["vision_model.conv1.weight"] = conv_w + # A text-decoder key the tower must leave untouched (and not route). + weights["model.layers.0.self_attn.q_proj.weight"] = torch.zeros(2, 2) + + tower.load_weights(weights) + self.assertTrue(torch.equal(tower.vit_large_projector.weight.detach(), proj_w)) + self.assertTrue(torch.equal(tower.vision_model.conv1.weight.detach(), conv_w)) + + +class TestStep3p7VLRegistration(unittest.TestCase): + """Verify the Step3p7 multimodal architecture and its encoder are registered.""" + + def test_vlm_entry_point_and_vision_encoder_registered(self): + from tensorrt_llm._torch.models.modeling_step3p7vl import ( + Step3p7VisionTower, + Step3p7VLForConditionalGeneration, + ) + from tensorrt_llm._torch.models.modeling_utils import ( + MODEL_CLASS_MAPPING, + MODEL_CLASS_VISION_ENCODER_MAPPING, + ) + + self.assertIs( + MODEL_CLASS_MAPPING.get("Step3p7ForConditionalGeneration"), + Step3p7VLForConditionalGeneration, + ) + entry = MODEL_CLASS_VISION_ENCODER_MAPPING.get("Step3p7ForConditionalGeneration") + self.assertIsNotNone(entry, "vision encoder not registered for Step3p7") + vision_cls, _vlm_base = entry + self.assertIs(vision_cls, Step3p7VisionTower) + + def test_input_processor_registered_for_model_type(self): + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3p7VLInputProcessor + + # ``register_input_processor`` stamps the model_type onto the processor + # class; the placeholder is the OOV-rewritten ```` token. + self.assertEqual(Step3p7VLInputProcessor._registered_model_type, "step3p7") + + +class TestStep3p7VLCheckpoint(unittest.TestCase): + """``vision_config`` geometry checks against the real Step-3.7-Flash checkpoints. + + The FP8 block-scale, NVFP4, and BF16 reference checkpoints all ship the same + PerceptionEncoder vision tower; only the text decoder's routed-expert + dtype/layout differs. The checkpoints are expected under ``LLM_MODELS_ROOT`` + (as in CI); a missing checkpoint surfaces as a test failure, not a skip. + """ + + def _check_vision_config(self, vision_cfg: dict): + self.assertEqual(vision_cfg["model_type"], "perception_encoder") + self.assertEqual(vision_cfg["width"], 1536) + self.assertEqual(vision_cfg["heads"], 16) + self.assertEqual(vision_cfg["layers"], 47) + self.assertEqual(vision_cfg["patch_size"], 14) + self.assertEqual(vision_cfg["image_size"], 728) + self.assertEqual(vision_cfg["hidden_act"], "quick_gelu") + self.assertEqual(vision_cfg["ls_init_value"], 0.1) + self.assertIs(vision_cfg["use_cls_token"], False) + # The vision MHA requires width divisible by the head count. + self.assertEqual(vision_cfg["width"] % vision_cfg["heads"], 0) + + @parameterized.expand( + [ + ("fp8", STEP3P7_FP8_DIR), + ("nvfp4", STEP3P7_NVFP4_DIR), + ("bf16", STEP3P7_BF16_DIR), + ] + ) + def test_vision_config_geometry(self, name, checkpoint_dir): + """All three checkpoints carry the same PerceptionEncoder vision config + plus the multimodal wiring (image token id, biasless projector).""" + config_dict = _load_config(checkpoint_dir) + self.assertIn("vision_config", config_dict) + self._check_vision_config(config_dict["vision_config"]) + self.assertEqual(config_dict["image_token_id"], 128001) + self.assertIs(config_dict.get("projector_bias", False), False) + + +class TestStep3p7VLInputProcessorHooks(unittest.TestCase): + """Multimodal-hashing hooks on ``Step3p7VLInputProcessor``. + + Step3 image spans interleave structural framing tokens + (````/````/````/````/ + ````) with the ```` embed slots. These tests + verify the processor exposes the per-image token count and framing-token + ids the generic hashing path needs to keep each image one contiguous span + (the KV-cache-reuse / chunked-prefill prerequisite). Requires the + Step-3.7-Flash checkpoint under ``LLM_MODELS_ROOT``. + """ + + @classmethod + def setUpClass(cls): + from transformers import AutoConfig, AutoTokenizer + + from tensorrt_llm._torch.models.modeling_step3p7vl import Step3p7VLInputProcessor + + cls.config = AutoConfig.from_pretrained(STEP3P7_BF16_DIR, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(STEP3P7_BF16_DIR, trust_remote_code=True) + cls.proc = Step3p7VLInputProcessor( + STEP3P7_BF16_DIR, cls.config, tokenizer, trust_remote_code=True + ) + + def test_special_and_mm_token_ids(self): + """All five framing tokens resolve; mm_token_ids = sentinel + framing.""" + special = self.proc.get_mm_special_token_ids() + self.assertIsNotNone(special) + self.assertEqual(special.numel(), 5) + # Distinct and resolved (not collapsed to a single unk id). + self.assertEqual(len(set(special.tolist())), 5) + + mm_ids = self.proc.get_mm_token_ids() + self.assertIsNotNone(mm_ids) + self.assertEqual(mm_ids[0].item(), self.proc._tllm_multimodal_token_id) + self.assertEqual(sorted(mm_ids[1:].tolist()), sorted(special.tolist())) + + def test_num_tokens_per_image_matches_processor(self): + """The hook delegates to the remote processor's span-length logic.""" + img = Image.new("RGB", (800, 600)) + n = self.proc.get_num_tokens_per_image(image=img) + self.assertEqual(n, self.proc._processor.get_num_image_tokens(800, 600)) + self.assertGreater(n, 0) + # CHW tensor (h, w) must agree with the PIL (w, h) path. + n_tensor = self.proc.get_num_tokens_per_image(image=torch.zeros(3, 600, 800)) + self.assertEqual(n_tensor, n) + + def test_image_span_is_contiguous(self): + """End-to-end: the hashing masks cover one contiguous image span whose + length equals ``get_num_tokens_per_image``.""" + from tensorrt_llm.inputs.multimodal import _compute_mm_masks + from tensorrt_llm.sampling_params import SamplingParams + + img = Image.new("RGB", (800, 600)) + inputs = {"prompt": "", "multi_modal_data": {"image": [img]}} + token_ids, _ = self.proc(inputs, SamplingParams()) + ids = torch.tensor(token_ids) + + mm_mask, embed_mask, special_mask = _compute_mm_masks( + ids, + vocab_size=self.proc.get_vocab_size(), + mm_token_ids=self.proc.get_mm_token_ids(), + mm_special_token_ids=self.proc.get_mm_special_token_ids(), + ) + + expected = self.proc.get_num_tokens_per_image(image=img) + self.assertEqual(int(mm_mask.sum()), expected) + # Embed slots (sentinels) + framing tokens partition the span. + self.assertEqual(int(embed_mask.sum()) + int(special_mask.sum()), expected) + # The span is a single contiguous run. + positions = mm_mask.nonzero().flatten() + self.assertEqual(int(positions[-1] - positions[0]) + 1, positions.numel()) + # Embed slots are exactly the OOV sentinels. + self.assertTrue(bool((ids[embed_mask] == self.proc._tllm_multimodal_token_id).all())) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 2eb25f381333..53fa03306f96 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -39,11 +39,11 @@ ExtendedRuntimePerfKnobConfig, KvCacheConfig, LookaheadDecodingConfig, MoeConfig, - PeftCacheConfig, PybindMirror, - RayPlacementConfig, SleepConfig, - SpeculativeConfig, StrictBaseModel, - TorchCompileConfig, TorchLlmArgs, - TrtLlmArgs, + MTPDecodingConfig, PeftCacheConfig, + PybindMirror, RayPlacementConfig, + SleepConfig, SpeculativeConfig, + StrictBaseModel, TorchCompileConfig, + TorchLlmArgs, TrtLlmArgs, UserProvidedDecodingConfig, update_llm_args_with_extra_dict) # fmt: on @@ -83,6 +83,21 @@ def test_LookaheadDecodingConfig(): assert pybind_config.max_verification_set_size == 4 +def test_MTPDecodingConfig_default_draft_len_is_not_user_set(): + config = MTPDecodingConfig() + + # Unset max_draft_len stays None (the "use the model's + # num_nextn_predict_layers" sentinel) until resolved at model load. + assert config.max_draft_len is None + assert config.max_total_draft_tokens is None + assert "max_draft_len" not in config.model_fields_set + + explicit_config = MTPDecodingConfig(max_draft_len=1) + assert explicit_config.max_draft_len == 1 + assert explicit_config.max_total_draft_tokens == 1 + assert "max_draft_len" in explicit_config.model_fields_set + + class TestYaml: def _yaml_to_dict(self, yaml_content: str) -> dict: