diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index af621d50caaa..f5cf71cdf584 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -425,7 +425,39 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, config.has_zero_point = layer_cfg['has_zero_point'] if 'pre_quant_scale' in layer_cfg: config.pre_quant_scale = layer_cfg['pre_quant_scale'] + # W4A16_NVFP4 is a modelopt label for full NVFP4 (W4A4). + # Normalize to NVFP4 for kernel dispatch so the CUTLASS + # NVFP4 MoE path is selected (the label distinction is + # only meaningful at the checkpoint-loading boundary). + if config.quant_algo == QuantAlgo.W4A16_NVFP4: + config.quant_algo = QuantAlgo.NVFP4 mixed_quant_configs[layer] = config + # Normalize "model.language_model." prefix to "model." so that + # quant_config_dict keys match TRT-LLM module names produced by + # named_modules() (which don't include the "language_model" level). + _LM_PREFIX = "model.language_model." + _MODEL_PREFIX = "model." + mixed_quant_configs = { + (_MODEL_PREFIX + k[len(_LM_PREFIX):] if k.startswith(_LM_PREFIX) else k): + v + for k, v in mixed_quant_configs.items() + } + # LMHead bypasses Linear.create_weights (manual Parameter), + # so NVFP4 weight scales are never allocated there. Move + # lm_head to exclude_modules so it loads as BF16. + lm_head_keys = [ + key for key in mixed_quant_configs + if key == "lm_head" or key.endswith(".lm_head") + ] + if lm_head_keys: + if quant_config.exclude_modules is None: + quant_config.exclude_modules = [] + quant_config.exclude_modules = list( + dict.fromkeys( + list(quant_config.exclude_modules) + ["lm_head"] + + lm_head_keys)) + for key in lm_head_keys: + del mixed_quant_configs[key] layer_quant_config = mixed_quant_configs elif quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES: if quant_config.group_size is None: diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py index bee5000939dc..8f203aba908a 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py @@ -1,3 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. import math import re from collections import defaultdict @@ -14,6 +28,24 @@ _FP8_2D_BLOCK_SIZE = 128 +# Suffixes that carry per-tensor scalar FP8 scales. These tensors have +# ndim == 0 (scalar shape ()) and must not be passed through any split/pack +# path that assumes a leading out-features dimension. Instead they are +# forwarded directly under the fused projection key because a single scalar +# applies uniformly across q, k, v, and z. +_SCALAR_SCALE_SUFFIXES = frozenset({"weight_scale", "input_scale"}) + +# NVFP4 (e2m1) lookup table: 16 representable values, index = nibble value. +# Bit layout: sign(1) | exponent(2) | mantissa(1). +# Values: 0, 0.5, 1, 1.5, 2, 3, 4, 6, -0, -0.5, -1, -1.5, -2, -3, -4, -6. +_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, +) + +# Block size for NVFP4 per-block weight scales (16 elements per block along K). +_NVFP4_BLOCK_SIZE = 16 + @register_mapper("HF", "Qwen3_5ForConditionalGeneration") @register_mapper("HF", "QwenImageBenchForConditionalGeneration") @@ -40,6 +72,9 @@ class Qwen3_5MoeHfWeightMapper(Qwen3NextHfWeightMapper): For FP8 checkpoints, the packed qkvz tensor is then dequantized to bf16 as a temporary workaround for TP loading (handled in _dequantize_linear_attn_fp8_qkvz). + Per-tensor scalar FP8 scales (weight_scale, input_scale, shape ()) + are forwarded directly to the fused qkvz key without splitting + (handled in _pack_split_projections). 3. MoE expert tensors (handled in handle_special_instance_module): Qwen3.5 BF16 checkpoints store fused gate_up_proj/down_proj per expert @@ -47,6 +82,15 @@ class Qwen3_5MoeHfWeightMapper(Qwen3NextHfWeightMapper): Qwen3.5 FP8 checkpoints store vanilla gate_proj/up_proj/down_proj per expert. This mapper detects which layout is present, transposes fused tensors, renames keys, and sets the matching MoEWeightLoadingMode. + + 4. NVFP4-quantized weights for excluded modules (handled in + _dequantize_nvfp4_excluded_weights): + Modules listed in quant_config.exclude_modules (e.g. lm_head) are + expected to load as BF16. However, the checkpoint may store their + weights in NVFP4 packed format (uint8, shape [N, K/2]) along with + per-block scales and a global scale. This method detects such tensors + by dtype and dequantizes them to BF16 before the module loader + attempts to copy them into the BF16 weight buffer. """ _SPLIT_PROJ_PATTERN = re.compile(r"^(.*\.linear_attn)\.in_proj_(qkv|q|k|v|z|b|a)\.(.+)$") @@ -55,6 +99,7 @@ class Qwen3_5MoeHfWeightMapper(Qwen3NextHfWeightMapper): ) def _normalize_weight_names(self, weights: dict) -> dict: + """Map HF checkpoint names onto the shared Qwen3Next module layout.""" normalized_weights = {} for key, tensor in weights.items(): if key.startswith("model.visual."): @@ -65,6 +110,12 @@ def _normalize_weight_names(self, weights: dict) -> dict: return normalized_weights def _normalize_scale_names(self, weights: dict, quant_algo) -> tuple[dict, bool]: + """Canonicalize ModelOpt FP8 scale names and shapes before loading. + + Returns: + The remapped weight dictionary and whether the source checkpoint + used ModelOpt's native FP8 block-scale layout. + """ # Canonicalize FP8 weight_scale layout so the Linear loader sees one # shape per quant algo: # - FP8_BLOCK_SCALES: modelopt fp8_pb_wo stores weight_scale shaped @@ -108,6 +159,7 @@ def handle_special_instance_module( module_weights: dict, allow_partial_loading: bool = False, ) -> None: + """Load Qwen3.5 MoE expert tensors with the layout TRT-LLM expects.""" if isinstance(module, MoE): config = self.config.pretrained_config uses_fused_expert_tensors = "gate_up_proj" in module_weights @@ -146,6 +198,7 @@ def handle_special_instance_module( ) def _pack_projection_tensor(self, tensors: list[torch.Tensor], num_groups: int) -> torch.Tensor: + """Group-interleave split projection tensors for fused linear attention.""" reference_shape = tensors[0].shape[1:] for tensor in tensors: assert tensor.shape[1:] == reference_shape, ( @@ -165,6 +218,7 @@ def _pack_projection_tensor(self, tensors: list[torch.Tensor], num_groups: int) def _split_qkv_tensor( self, tensor: torch.Tensor, expected_q: int, expected_v: int ) -> tuple[torch.Tensor, ...]: + """Split a packed qkv tensor into q, k, and v component tensors.""" expected_total = expected_q * 2 + expected_v assert tensor.shape[0] == expected_total, ( f"Expected packed qkv projection with leading dim {expected_total}, got {tensor.shape}" @@ -174,6 +228,7 @@ def _split_qkv_tensor( def _split_qkv_scale_tensor( self, tensor: torch.Tensor, expected_q: int, expected_v: int ) -> tuple[torch.Tensor, ...]: + """Split a packed qkv FP8 block-scale tensor into q, k, and v scales.""" expected_q_blocks = math.ceil(expected_q / _FP8_2D_BLOCK_SIZE) expected_v_blocks = math.ceil(expected_v / _FP8_2D_BLOCK_SIZE) expected_total_blocks = expected_q_blocks * 2 + expected_v_blocks @@ -186,6 +241,7 @@ def _split_qkv_scale_tensor( def _dequantize_fp8_block_scale_weight( self, weight: torch.Tensor, weight_scale_inv: torch.Tensor ) -> torch.Tensor: + """Dequantize a 2-D FP8 block-scale linear-attention weight tensor.""" rows, cols = weight.shape expanded_scales = ( weight_scale_inv.to(torch.float32) @@ -198,6 +254,7 @@ def _dequantize_fp8_block_scale_weight( return (weight.to(torch.float32) * expanded_scales).to(target_dtype).contiguous() def _dequantize_linear_attn_fp8_qkvz(self, weights: dict) -> dict: + """Dequantize packed qkvz FP8 block-scale weights that load as BF16.""" updated_weights = dict(weights) for name in list(weights): if not name.endswith(".linear_attn.in_proj_qkvz.weight"): @@ -211,7 +268,180 @@ def _dequantize_linear_attn_fp8_qkvz(self, weights: dict) -> dict: updated_weights.pop(scale_name, None) return updated_weights + def _dequantize_nvfp4_weight( + self, + weight_uint8: torch.Tensor, + block_scale_fp8: torch.Tensor, + global_scale_fp32: torch.Tensor, + ) -> torch.Tensor: + """Dequantize a 2-D NVFP4 weight tensor to BF16. + + ``weight_uint8`` is ``(N, K/2)`` packed (two e2m1 nibbles per byte), + ``block_scale_fp8`` is ``(N, K/16)`` (fp8_e4m3, block size 16 along K), + ``global_scale_fp32`` is a scalar float32 tensor. + + Returns a ``(N, K)`` BF16 tensor. + """ + lut = _E2M1_VALUES.to(weight_uint8.device) + N, K_half = weight_uint8.shape + K = K_half * 2 + + low = weight_uint8 & 0x0F + high = (weight_uint8 >> 4) & 0x0F + + vals = torch.empty(N, K, dtype=torch.float32, device=weight_uint8.device) + vals[:, 0::2] = lut[low.long()] + vals[:, 1::2] = lut[high.long()] + + # block_scale_fp8: (N, K/16) — each scale covers 16 K elements + scale = (block_scale_fp8.to(torch.float32) * global_scale_fp32.to(torch.float32)).unsqueeze( + -1 + ) # (N, K/16, 1) + + vals = vals.view(N, K // _NVFP4_BLOCK_SIZE, _NVFP4_BLOCK_SIZE) * scale + target_dtype = getattr(self.config.pretrained_config, "torch_dtype", torch.bfloat16) + if target_dtype is None: + target_dtype = torch.bfloat16 + return vals.view(N, K).to(target_dtype).contiguous() + + def _dequantize_nvfp4_excluded_weights(self, weights: dict) -> dict: + """Dequantize NVFP4-packed weights for modules in exclude_modules. + + Modules listed in quant_config.exclude_modules (e.g. lm_head) should + load as BF16, but NVFP4 checkpoints store their weights as packed uint8 + with per-block fp8 scales and a global fp32 scale. Detect these by + checking for uint8 weight tensors alongside the matching scale tensors, + and dequantize to BF16 in-place in the weight dict. + + Scale tensors (weight_scale, weight_scale_2, input_scale) for the + dequantized modules are also removed from the dict because the module's + BF16 weight buffer has no slots for them. + """ + qc = self.config.quant_config + if qc is None or qc.exclude_modules is None: + return weights + + updated = dict(weights) + # Collect all weight keys for excluded modules with uint8 dtype. + # Key format after normalization: ".weight" + candidates = [ + k for k, v in weights.items() if k.endswith(".weight") and v.dtype == torch.uint8 + ] + for weight_key in candidates: + prefix = weight_key[: -len(".weight")] + # Check if this prefix corresponds to an excluded module. + if not qc.is_module_excluded_from_quantization(prefix): + continue + block_scale_key = f"{prefix}.weight_scale" + global_scale_key = f"{prefix}.weight_scale_2" + if block_scale_key not in weights or global_scale_key not in weights: + # Not a standard NVFP4 layout; skip. + continue + updated[weight_key] = self._dequantize_nvfp4_weight( + weights[weight_key], + weights[block_scale_key], + weights[global_scale_key], + ) + # Remove scale tensors — the BF16 module has no parameter slots for them. + for scale_suffix in ("weight_scale", "weight_scale_2", "input_scale"): + scale_key = f"{prefix}.{scale_suffix}" + updated.pop(scale_key, None) + + return updated + + def _dequantize_fp8_pertensor_excluded_split_weights(self, weights: dict) -> dict: + """Dequantize FP8 per-tensor-scale split linear-attention projections. + + In MIXED_PRECISION checkpoints (e.g. Qwen3.6-35B-A3B-NVFP4) the + linear-attention projections are stored as split FP8 tensors: + in_proj_qkv.weight (float8_e4m3fn, per-tensor scalar weight_scale) + in_proj_z.weight (float8_e4m3fn, per-tensor scalar weight_scale) + These are packed into ``in_proj_qkvz`` by _pack_split_projections. + + The global quant_algo is MIXED_PRECISION, but TRT-LLM does not have a + per-layer FP8 quant entry for the packed ``in_proj_qkvz`` name. As a + result the in_proj_qkvz Linear is constructed with UnquantizedLinearMethod + (BF16 weight buffer). Without this method, ``copy_weight`` performs a raw + dtype cast from FP8 to BF16, interpreting FP8 bit-patterns directly — + inflating each value by ~1/weight_scale (≈ 1000× for typical scales). + + This method runs BEFORE _pack_split_projections. For every split FP8 + linear-attention projection weight that carries a per-tensor scalar + weight_scale, it dequantizes the component individually via: + bf16 = fp8.to(float32) * weight_scale + then removes the associated scale tensors (weight_scale, input_scale) so + _pack_split_projections never emits a now-wrong fused scalar scale for + the packed BF16 module. + + Only scalar (ndim == 0) per-tensor scales are handled here; 2D-block + FP8 scales (from FP8_BLOCK_SCALES checkpoints) are handled by + _dequantize_linear_attn_fp8_qkvz which runs after packing. + + Guard: only applies to MIXED_PRECISION quant_algo because that is the + only mode where ``in_proj_qkvz`` receives an unquantized (BF16) Linear + buffer despite FP8 source weights. For FP8_BLOCK_SCALES or FP8_QDQ + the weight buffer IS FP8-typed so the raw copy is correct. + """ + _FP8_DTYPES = (torch.float8_e4m3fn,) + try: + _FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e4m3fnuz) + except AttributeError: + pass + + target_dtype = ( + getattr(self.config.pretrained_config, "torch_dtype", torch.bfloat16) or torch.bfloat16 + ) + + updated = dict(weights) + # Walk all FP8 weight keys and find split linear-attention projections + # (matching the pattern .linear_attn.in_proj_{qkv|q|k|v|z}.weight) + # that carry a scalar per-tensor weight_scale. + candidates = [ + k for k, v in weights.items() if k.endswith(".weight") and v.dtype in _FP8_DTYPES + ] + for weight_key in candidates: + # Match the full weight_key, e.g. + # "model.layers.0.linear_attn.in_proj_qkv.weight" + # Group 3 of the pattern captures the suffix ("weight" here). + match = self._SPLIT_PROJ_PATTERN.match(weight_key) + if match is None: + continue + _attn_prefix, proj_name, matched_suffix = match.groups() + if matched_suffix != "weight": + continue # skip scale tensors caught by the same pattern + if proj_name not in {"qkv", "q", "k", "v", "z", "b", "a"}: + continue + # Only dequantize projections that will be packed into in_proj_qkvz + # or in_proj_ba (both receive BF16 buffers under MIXED_PRECISION). + # Check for a scalar per-tensor weight_scale on this split tensor. + scale_key = weight_key[: -len(".weight")] + ".weight_scale" + if scale_key not in weights: + continue + scale = weights[scale_key] + if scale.ndim != 0: + continue # 2D-block scale — handled by _dequantize_linear_attn_fp8_qkvz + # Dequantize: bf16 = fp8.to(float32) * weight_scale + updated[weight_key] = ( + (weights[weight_key].to(torch.float32) * scale.to(torch.float32)) + .to(target_dtype) + .contiguous() + ) + # Remove scale tensors — the packed BF16 module has no parameter + # slots for them, and _pack_split_projections must not emit a + # now-stale fused scalar scale under the in_proj_qkvz key. + for scale_suffix in ("weight_scale", "input_scale"): + scale_k = weight_key[: -len(".weight")] + f".{scale_suffix}" + updated.pop(scale_k, None) + + return updated + def _pack_split_projections(self, weights: dict) -> dict: + """Pack Qwen3.5 split linear-attention projections into fused keys. + + Qwen3.5 checkpoints may store qkv/z or q/k/v/z plus b/a separately. + TRT-LLM's Qwen3Next modules expect fused ``in_proj_qkvz`` and + ``in_proj_ba`` tensors with grouped interleaving. + """ config = self.config.pretrained_config num_k_groups = config.linear_num_key_heads num_v_heads = config.linear_num_value_heads @@ -236,6 +466,46 @@ def _pack_split_projections(self, weights: dict) -> dict: expected_ba = config.linear_num_value_heads for (prefix, suffix), tensors in grouped_weights.items(): + # Per-tensor scalar FP8 scales (weight_scale, input_scale) have + # ndim == 0, i.e. shape (). They cannot be split along an + # out-features axis because they are a single scalar that applies + # uniformly to the entire projection. Forward the scalar from the + # qkv sub-tensor directly as the fused qkvz/ba key and skip the + # regular split/pack path entirely. + if suffix in _SCALAR_SCALE_SUFFIXES and all(t.ndim == 0 for t in tensors.values()): + qkvz_candidates = {"qkv", "q", "k", "v", "z"} & tensors.keys() + if qkvz_candidates: + # Use the scalar from "qkv" if present, else fall back to + # any available candidate (q/k/v all hold the same value). + representative = ( + tensors["qkv"] + if "qkv" in tensors + else next(v for k, v in tensors.items() if k in {"q", "k", "v"}) + ) + assert representative.ndim == 0, ( + f"Expected scalar (ndim=0) for {prefix}.in_proj_qkv.{suffix}, " + f"got shape {representative.shape}" + ) + packed_name = f"{prefix}.in_proj_qkvz.{suffix}" + assert packed_name not in packed_weights, ( + f"Packed projection {packed_name} already exists" + ) + packed_weights[packed_name] = representative + + ba_candidates = {"b", "a"} & tensors.keys() + if ba_candidates: + representative_ba = tensors["b"] if "b" in tensors else tensors["a"] + assert representative_ba.ndim == 0, ( + f"Expected scalar (ndim=0) for {prefix}.in_proj_b.{suffix}, " + f"got shape {representative_ba.shape}" + ) + packed_name = f"{prefix}.in_proj_ba.{suffix}" + assert packed_name not in packed_weights, ( + f"Packed projection {packed_name} already exists" + ) + packed_weights[packed_name] = representative_ba + continue + # `weight_scale_inv` (loaded by FP8BlockScalesLinearMethod) is the # only suffix stored in 2D-block format [ceil(out/block_size), ceil(in/block_size)]; # weight/bias/weight_scale all keep out_features as their leading dim. @@ -317,6 +587,7 @@ def _remap_dense_mlp_weights(self, weights: dict) -> dict: return remapped_weights def preprocess_weights(self, weights: dict) -> dict: + """Normalize, pack, and dequantize Qwen3.5 weights before loading.""" quant_algo = self.config.quant_config.quant_algo normalized_weights = self._normalize_weight_names(weights) @@ -324,10 +595,27 @@ def preprocess_weights(self, weights: dict) -> dict: normalized_weights, quant_algo ) + # For MIXED_PRECISION checkpoints with FP8 per-tensor-scale linear-attn + # projections (e.g. Qwen3.6-35B-A3B-NVFP4): dequantize the split FP8 + # weights BEFORE packing so each component uses its own scale. + # Must run before _pack_split_projections. + if quant_algo == QuantAlgo.MIXED_PRECISION: + normalized_weights = self._dequantize_fp8_pertensor_excluded_split_weights( + normalized_weights + ) + packed_weights = self._pack_split_projections(normalized_weights) if quant_algo == QuantAlgo.FP8_BLOCK_SCALES and not is_modelopt_pb_wo: packed_weights = self._dequantize_linear_attn_fp8_qkvz(packed_weights) + # For MIXED_PRECISION checkpoints, modules in exclude_modules (e.g. + # lm_head) should load as BF16 but the checkpoint may store them as + # NVFP4 (uint8 weight + fp8 block scales + fp32 global scale). + # Dequantize those weights here so the BF16 module can copy them + # directly without a dtype/shape mismatch. + if quant_algo == QuantAlgo.MIXED_PRECISION: + packed_weights = self._dequantize_nvfp4_excluded_weights(packed_weights) + if not getattr(self.config.pretrained_config, "num_experts", 0): packed_weights = self._remap_dense_mlp_weights(packed_weights) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 7667972804ad..4e2bc6bda341 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -36,6 +36,7 @@ from tensorrt_llm._torch.pyexecutor.config_utils import \ get_qwen3_hybrid_layer_types from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.quantization import QuantAlgo from ...logger import logger from ..attention_backend import AttentionMetadata @@ -146,6 +147,24 @@ def __init__( weight_loading_mode = (MoEWeightLoadingMode.FUSED_GATE_UP_PROJ if config.model_type == "qwen3_5_moe_text" else MoEWeightLoadingMode.VANILLA) + # For MIXED_PRECISION checkpoints (e.g. Qwen3.6-35B-A3B-NVFP4) the + # global quant_algo is MIXED_PRECISION but each MoE layer has a + # per-layer NVFP4 quant config. Pass it explicitly so create_moe() + # selects the right MoE kernel (NVFP4-aware) instead of falling back + # to the default CutlassFusedMoE which cannot load NVFP4 weights. + moe_override_quant_config = None + if (model_config.quant_config_dict is not None + and model_config.quant_config.quant_algo + == QuantAlgo.MIXED_PRECISION and layer_idx is not None): + candidate_keys = [ + f"model.language_model.layers.{layer_idx}.mlp.experts", + f"model.layers.{layer_idx}.mlp.experts", + ] + for key in candidate_keys: + if key in model_config.quant_config_dict: + moe_override_quant_config = model_config.quant_config_dict[ + key] + break self.experts = create_moe( num_experts=self.num_experts, routing_method=self.gate.routing_method, @@ -157,6 +176,7 @@ def __init__( model_config=model_config, layer_idx=layer_idx, weight_loading_mode=weight_loading_mode, + override_quant_config=moe_override_quant_config, ) self.shared_expert = GatedMLP( diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 0bd2f8400157..9a5f3a8851e2 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1451,7 +1451,7 @@ def __init__( case "nemotron_h" | "nemotron_h_puzzle": from .modeling_nemotron_h import NemotronHMTP mtp_layer = NemotronHMTP - case "qwen3_next" | "qwen3_5_text" | "qwen3_5_moe_text": + case "qwen3_next" | "qwen3_5_text" | "qwen3_5_moe" | "qwen3_5_moe_text": from .modeling_qwen3_next import Qwen3NextMTP mtp_layer = Qwen3NextMTP case "step3p7" | "step3p5": @@ -1513,7 +1513,9 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], layer_idx, aux_stream_dict, is_separate_draft_engine=False) - elif model_type == "qwen3_next": + elif model_type in [ + "qwen3_next", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text" + ]: from .modeling_qwen3_next import Qwen3NextMTP mtp_layer = Qwen3NextMTP(model_config, layer_idx, aux_stream_dict) else: diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index e949bf4eb754..bb25b9685f38 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -475,8 +475,24 @@ def _flatten_rope(text_config: dict) -> dict: has_mrope = ("mrope_section" in rope_scaling or rope_scaling.get("mrope_interleaved", False)) if has_mrope: - rope_scaling["type"] = "mrope" + # Qwen3.5 VLM checkpoints embed mrope_section / mrope_interleaved + # in rope_parameters for use with the vision path. The text + # executor never constructs 3D position_ids (no vision encoder), + # so leaving type="mrope" causes MRotaryEmbedding to silently + # produce wrong cos/sin from 2D position_ids. Strip the mRoPE + # fields unconditionally here; partial_rotary_factor (already + # extracted above) carries the fractional-RoPE scaling needed + # for the linear-attention layers. + rope_scaling.pop("mrope_section", None) + rope_scaling.pop("mrope_interleaved", None) rope_scaling.pop("rope_type", None) + rope_scaling.pop("type", None) + # After stripping the mRoPE fields, what remains (if anything) + # is standard scaling config. If nothing meaningful is left, + # clear rope_scaling to avoid triggering unexpected code paths. + if rope_scaling: + if "type" not in rope_scaling and "rope_type" not in rope_scaling: + rope_scaling = {} elif "type" not in rope_scaling and "rope_type" in rope_scaling: rope_type = rope_scaling.pop("rope_type") # "default" means standard RoPE (no scaling) — don't set @@ -487,6 +503,10 @@ def _flatten_rope(text_config: dict) -> dict: rope_scaling["type"] = rope_type if rope_scaling: text_config["rope_scaling"] = rope_scaling + else: + # Clearing rope_scaling locally is not enough — the original key + # in text_config still points at the pre-strip dict. Remove it. + text_config.pop("rope_scaling", None) return text_config diff --git a/tensorrt_llm/quantization/mode.py b/tensorrt_llm/quantization/mode.py index e4e8fbc89d17..d1f3b0a852bc 100644 --- a/tensorrt_llm/quantization/mode.py +++ b/tensorrt_llm/quantization/mode.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -40,6 +40,10 @@ class QuantAlgo(StrEnum, metaclass=BaseEnumMeta): INT8 = auto() MIXED_PRECISION = auto() NVFP4 = auto() + # W4A16_NVFP4 is a modelopt naming convention alias for NVFP4. + # Both weight and input activations are quantized to FP4 (W4A4). + # The kernel dispatch path normalizes this to NVFP4 semantics. + W4A16_NVFP4 = auto() W4A8_NVFP4_FP8 = auto() W4A8_MXFP4_FP8 = auto() W4A8_MXFP4_MXFP8 = auto() @@ -418,6 +422,10 @@ def from_quant_algo( elif quant_algo == QuantAlgo.NVFP4_ARC: # NVFP4_ARC uses the same QuantMode as NVFP4, distinction is at QuantAlgo level quant_mode = QuantMode.from_description(use_nvfp4=True) + elif quant_algo == QuantAlgo.W4A16_NVFP4: + # W4A16_NVFP4 is a modelopt label for NVFP4 (full W4A4 FP4 quant). + # Map to the same QuantMode bits as NVFP4 for kernel dispatch. + quant_mode = QuantMode.from_description(use_nvfp4=True) elif quant_algo == QuantAlgo.W4A8_NVFP4_FP8: quant_mode = QuantMode.from_description(use_w4a8_nvfp4_fp8=True) elif quant_algo == QuantAlgo.W4A8_MXFP4_FP8: diff --git a/tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py b/tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py new file mode 100644 index 000000000000..82da0a0e708b --- /dev/null +++ b/tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. +"""Unit tests for Qwen3.5/Qwen3.6 checkpoint weight mapping.""" + +import types + +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +def _make_mapper() -> Qwen3_5MoeHfWeightMapper: + pretrained_config = types.SimpleNamespace( + linear_key_head_dim=2, + linear_value_head_dim=2, + linear_num_key_heads=1, + linear_num_value_heads=1, + num_hidden_layers=1, + num_experts=1, + torch_dtype=torch.bfloat16, + ) + model_config = ModelConfig( + pretrained_config=pretrained_config, + mapping=Mapping(), + quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION), + ) + mapper = object.__new__(Qwen3_5MoeHfWeightMapper) + mapper._config = model_config + return mapper + + +def test_fp8_pertensor_linear_attn_weights_are_dequantized_before_pack(): + mapper = _make_mapper() + scale = torch.tensor(0.25, dtype=torch.float32) + qkv_fp8 = torch.tensor( + [ + [1.0, -2.0], + [3.0, -4.0], + [0.5, -0.5], + [1.5, -1.5], + [2.0, -1.0], + [4.0, -3.0], + ], + dtype=torch.float8_e4m3fn, + ) + z_fp8 = torch.tensor([[2.0, -1.0], [1.0, -0.5]], dtype=torch.float8_e4m3fn) + weights = { + "model.layers.0.linear_attn.in_proj_qkv.weight": qkv_fp8, + "model.layers.0.linear_attn.in_proj_qkv.weight_scale": scale, + "model.layers.0.linear_attn.in_proj_qkv.input_scale": torch.tensor( + 1.0, dtype=torch.float32 + ), + "model.layers.0.linear_attn.in_proj_z.weight": z_fp8, + "model.layers.0.linear_attn.in_proj_z.weight_scale": scale, + "model.layers.0.linear_attn.in_proj_z.input_scale": torch.tensor(1.0, dtype=torch.float32), + } + + packed = mapper.preprocess_weights(weights) + + packed_weight = packed["model.layers.0.linear_attn.in_proj_qkvz.weight"] + expected = torch.cat( + [ + qkv_fp8[0:2].to(torch.float32) * scale, + qkv_fp8[2:4].to(torch.float32) * scale, + qkv_fp8[4:6].to(torch.float32) * scale, + z_fp8.to(torch.float32) * scale, + ], + dim=0, + ).to(torch.bfloat16) + assert packed_weight.dtype == torch.bfloat16 + torch.testing.assert_close(packed_weight, expected) + assert "model.layers.0.linear_attn.in_proj_qkvz.weight_scale" not in packed + assert "model.layers.0.linear_attn.in_proj_qkvz.input_scale" not in packed diff --git a/tests/unittest/llmapi/test_llm_quant.py b/tests/unittest/llmapi/test_llm_quant.py index 642571814267..20dfacb73ab4 100644 --- a/tests/unittest/llmapi/test_llm_quant.py +++ b/tests/unittest/llmapi/test_llm_quant.py @@ -216,6 +216,50 @@ def test_quant_cfg_from_hf_quant_config(): assert layer_quant_config["model.layers.0.mlp.up_proj"].group_size == 64 +def test_quant_cfg_qwen35_nvfp4_alias_and_prefix_normalization(): + """Qwen3.5/3.6 MIXED_PRECISION NVFP4 layers use TRT-LLM module names.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = model_dir / "hf_quant_config.json" + with open(hf_quant_config_file, 'w') as f: + json.dump( + { + "producer": { + "name": "modelopt" + }, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "FP8", + "quantized_layers": { + "model.language_model.layers.0.mlp.experts": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, + "lm_head": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, + "model.language_model.lm_head": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, + }, + }, + }, f) + + quant_config, layer_quant_config = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, "CUTLASS") + + assert quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + assert quant_config.exclude_modules == ["lm_head", "model.lm_head"] + assert "lm_head" not in layer_quant_config + assert "model.lm_head" not in layer_quant_config + assert "model.language_model.layers.0.mlp.experts" not in layer_quant_config + experts_config = layer_quant_config["model.layers.0.mlp.experts"] + assert experts_config.quant_algo == QuantAlgo.NVFP4 + assert experts_config.group_size == 16 + + def _write_hf_quant_config(model_dir: Path, content: dict) -> Path: """Write a ``hf_quant_config.json`` under ``model_dir`` and return its path.""" path = model_dir / "hf_quant_config.json"