From 65020585b4dc11a01a2b976b6fd467176bfd2d6f Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Wed, 24 Jun 2026 14:06:06 -0700 Subject: [PATCH 01/13] [None][feat] Port Qwen3.6-35B-A3B-NVFP4 + MTP to TRT-LLM PyTorch backend Four changes required to load nvidia/Qwen3.6-35B-A3B-NVFP4 on SM121: 1. quantization/mode.py: Add W4A16_NVFP4 to QuantAlgo enum as an alias for NVFP4. modelopt labels the Qwen3.6 NVFP4 checkpoint with "W4A16_NVFP4" (a naming convention artifact -- both weight and activations are quantized to FP4 / W4A4). Map it to the same QuantMode bits as NVFP4 so the CUTLASS kernel dispatch is unchanged. 2. _torch/model_config.py: In _build_modelopt_quant_config, normalize per-layer W4A16_NVFP4 -> NVFP4 for kernel dispatch and move lm_head to exclude_modules so it loads as BF16 (LMHead bypasses Linear.create_weights and cannot carry NVFP4 scale Parameters). 3. _torch/pyexecutor/config_utils.py: Strip mrope_section and mrope_interleaved from _Qwen35ConfigCompat._flatten_rope for the text executor path. The Qwen3.6 VLM checkpoint includes mRoPE fields intended for the vision path; the text executor never builds 3D position_ids so leaving type="mrope" causes MRotaryEmbedding to produce silently wrong cos/sin from 2D position_ids. 4. _torch/models/modeling_speculative.py: Extend MTPDraftModel (two-engine mode) to dispatch qwen3_5_text and qwen3_5_moe_text to Qwen3NextMTP. The one-engine path (MTPOneDraftModel) already included these model types; this brings the two-engine path into parity. The MTP weight key normalization (mtp.* -> model.layers.40.*), the NVFP4 MoE scale loading, and the MTP unquantized-sublayer backend fallback are all already handled by existing code in Qwen3NextHfWeightMapper and Qwen3NextMTP. Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/model_config.py | 16 ++++++++++++++++ .../_torch/models/modeling_speculative.py | 2 +- tensorrt_llm/_torch/pyexecutor/config_utils.py | 17 ++++++++++++++++- tensorrt_llm/quantization/mode.py | 10 +++++++++- 4 files changed, 42 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index af621d50caaa..b5b2ab06b91f 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -425,7 +425,23 @@ 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 + # 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. + if mixed_quant_configs and "lm_head" in mixed_quant_configs: + if quant_config.exclude_modules is None: + quant_config.exclude_modules = [] + if "lm_head" not in quant_config.exclude_modules: + quant_config.exclude_modules = list( + quant_config.exclude_modules) + ["lm_head"] + del mixed_quant_configs["lm_head"] 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/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 0bd2f8400157..cdfcbd68ba86 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1513,7 +1513,7 @@ 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_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..bf3ba5b20ede 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -475,8 +475,23 @@ 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) + # 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 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: From 1a6cc2bb91fa44d3a5f80343a5ecfc1f0225ee3e Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Wed, 24 Jun 2026 17:08:09 -0700 Subject: [PATCH 02/13] [None][fix] Remove stale rope_scaling key after mRoPE strip The mRoPE strip block worked on a local copy of rope_scaling but only wrote back when the dict was non-empty. When the strip cleared the dict, text_config["rope_scaling"] retained the pre-strip mrope_section / mrope_interleaved values, silently defeating the strip and producing wrong cos/sin from MRotaryEmbeddings 2D fallback at inference time. Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/pyexecutor/config_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index bf3ba5b20ede..795f2f065243 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -502,6 +502,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 From 4e801400848384746cd84ffbf76c762a2eb6e8fd Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Fri, 26 Jun 2026 09:48:02 -0700 Subject: [PATCH 03/13] [None][fix] qwen3_5_weight_mapper: FP8 scalar pass-through + lm_head NVFP4 dequant Scalar FP8 scale tensors (weight_scale, input_scale) have ndim == 0 and cannot be split or stacked across q/k/v/z projections. Forward them directly to the fused qkvz/ba key instead of routing through the regular split/pack path (_SCALAR_SCALE_SUFFIXES early-exit in _pack_split_projections). Add _dequantize_nvfp4_weight + _dequantize_nvfp4_excluded_weights: modules in quant_config.exclude_modules (e.g. lm_head) are expected to load as BF16, but NVFP4 checkpoints store them as packed uint8 with per-block fp8 scales and a global fp32 scale. Detect by dtype, dequant to BF16, and drop the orphaned scale tensors before the module loader runs. Call site added in preprocess_weights for MIXED_PRECISION quant. Also add NVIDIA copyright header (file previously had none). Signed-off-by: Mihai Chiorean --- .../checkpoints/hf/qwen3_5_weight_mapper.py | 176 +++++++++++++++++- 1 file changed, 175 insertions(+), 1 deletion(-) 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..b3983985c71e 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,25 @@ _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 +73,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 +83,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)\.(.+)$") @@ -60,7 +105,7 @@ def _normalize_weight_names(self, weights: dict) -> dict: if key.startswith("model.visual."): continue if key.startswith("model.language_model."): - key = "model." + key[len("model.language_model.") :] + key = "model." + key[len("model.language_model."):] normalized_weights[key] = tensor return normalized_weights @@ -211,6 +256,89 @@ 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 _pack_split_projections(self, weights: dict) -> dict: config = self.config.pretrained_config num_k_groups = config.linear_num_key_heads @@ -236,6 +364,44 @@ 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: + 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.get("qkv") or 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.get("b") or tensors.get("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. @@ -328,6 +494,14 @@ def preprocess_weights(self, weights: dict) -> dict: 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) From 875f2123caac88413bd5cf5d6ae1acef41a718f1 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Fri, 26 Jun 2026 13:50:29 -0700 Subject: [PATCH 04/13] [None][fix] Qwen3.6 NVFP4 MoE: normalize lang_model prefix + per-layer dispatch Two correctness fixes for loading Qwen3.6-35B-A3B-NVFP4 on the PyTorch backend: 1. model_config.py: hf_quant_config.json keys use the VLM-wrapper prefix "model.language_model.<...>", but TRT-LLM module names produced by named_modules() (after the _Qwen35ConfigCompat shim strips the vision wrapper) use "model.<...>". Normalize keys in _build_modelopt_quant_config() so apply_layerwise_quant_config() actually matches. Without this, every Linear/MoE module sees the global MIXED_PRECISION config and ignores its per-layer NVFP4 quant_algo. 2. modeling_qwen3_next.py: Qwen3NextSparseMoeBlock.__init__ now looks up the per-layer MoE NVFP4 config and passes it as override_quant_config to create_moe(). Without this, get_moe_cls() sees MIXED_PRECISION and falls back to CutlassFusedMoE (BF16 expectations), which then fails with a 2048 vs 512 NVFP4 tensor shape mismatch when loading expert weights. Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/model_config.py | 9 +++++++++ .../_torch/models/modeling_qwen3_next.py | 19 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index b5b2ab06b91f..fc2451d70be5 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -432,6 +432,15 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, 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. diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 7667972804ad..322e35600d0e 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -42,6 +42,7 @@ from ..distributed import (AllReduce, AllReduceFusionOp, AllReduceParams, MoEAllReduce, MoEAllReduceParams, allgather) from ..model_config import ModelConfig +from tensorrt_llm.quantization import QuantAlgo from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.fused_moe import (BaseMoeRoutingMethod, MoEWeightLoadingMode, @@ -146,6 +147,23 @@ 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 +175,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( From b1f015089300207da34020376d9dd362c1eb2f3d Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Fri, 26 Jun 2026 15:28:59 -0700 Subject: [PATCH 05/13] [None][fix] dequantize FP8 per-tensor linear-attn weights before packing in NVFP4 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Qwen3.6-35B-A3B-NVFP4 stores linear-attention in_proj_qkv and in_proj_z as FP8 per-tensor-scale tensors (float8_e4m3fn + scalar weight_scale). These are packed into in_proj_qkvz by _pack_split_projections. With a MIXED_PRECISION global quant_algo, the in_proj_qkvz Linear receives UnquantizedLinearMethod (BF16 weight buffer) because there is no matching per-layer FP8 entry under the packed TRT-LLM module name. Previously, copy_weight cast the packed FP8 tensor to BF16 numerically — interpreting raw FP8 bit-patterns without applying the per-tensor weight_scale. This inflated each weight value by ~1/weight_scale (~1000x), corrupting the linear-attention hidden states in all 30 GatedDeltaNet layers. Fix: add _dequantize_fp8_pertensor_excluded_split_weights, called in preprocess_weights before _pack_split_projections for MIXED_PRECISION. It detects split linear-attention FP8 weight tensors with a scalar weight_scale and dequantizes each component individually: bf16 = fp8.to(float32) * weight_scale Orphaned scale tensors (weight_scale, input_scale) are removed so _pack_split_projections does not forward a now-stale scalar scale. Smoke test on Qwen3.6-35B-A3B-NVFP4: raw: "The capital of France is" -> " Paris, a city renowned..." chat: "What is the capital of France?" -> coherent English reasoning Signed-off-by: Mihai Chiorean --- .../checkpoints/hf/qwen3_5_weight_mapper.py | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) 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 b3983985c71e..205fd0c87af4 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 @@ -339,6 +339,92 @@ def _dequantize_nvfp4_excluded_weights(self, weights: dict) -> dict: 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: config = self.config.pretrained_config num_k_groups = config.linear_num_key_heads @@ -490,6 +576,15 @@ 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) From 0c5938d9ee939a287fd7bb57b664d6ed0b4e6202 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Fri, 26 Jun 2026 15:52:04 -0700 Subject: [PATCH 06/13] [None][fix] Guard scalar-scale bypass dim + pop legacy mrope type Two defensive fixes from codex review: 1. qwen3_5_weight_mapper: the scalar weight_scale / input_scale bypass for FP8 attention projections was firing for any matching suffix regardless of dimensionality. FP8_PER_CHANNEL_PER_TOKEN checkpoints have 1-D weight_scale tensors (after _normalize_scale_names squeezes the trailing axis), which would trip the ndim==0 assertion and abort weight loading. Require all candidate tensors to be 0-dim before taking the bypass so per-channel scales fall through to the regular split/pack path. 2. _Qwen35ConfigCompat._flatten_rope: also pop the legacy "type" key alongside "rope_type" / "mrope_section" / "mrope_interleaved". HF Qwen2.5-VL configs used {"type": "mrope", ...}; the new rope_type form (HF transformers 5.x) is what Qwen3.5/3.6 ship, but defending against the legacy form is cheap and avoids resurrecting the broken MRotaryEmbedding path on older configs. Signed-off-by: Mihai Chiorean --- .../_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py | 3 ++- tensorrt_llm/_torch/pyexecutor/config_utils.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) 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 205fd0c87af4..c669d7dfc3be 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 @@ -456,7 +456,8 @@ def _pack_split_projections(self, weights: dict) -> dict: # 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: + 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 diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 795f2f065243..bb25b9685f38 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -486,6 +486,7 @@ def _flatten_rope(text_config: dict) -> dict: 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. From 1618f82e5311ad413b9ca795ade52ebd03f05e05 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Fri, 26 Jun 2026 19:51:36 -0700 Subject: [PATCH 07/13] [None][fix] MTP dispatch: also accept qwen3_5_moe (VLM-composite) The text-only path normalizes to model_type=qwen3_5_moe_text and is exercised by the smoke test. The VLM-composite checkpoint format publishes model_type=qwen3_5_moe at the top level (before the _Qwen35ConfigCompat shim strips the wrapper), which would otherwise fall through to the unsupported-model-type error on the MTP draft engine path. Both MTPForCausalLM (one-model) and MTPDraftModel (separate engine) dispatches now cover the composite form for forward compatibility. Codex pre-PR review flag. Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/models/modeling_speculative.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index cdfcbd68ba86..4867ecf68661 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,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], layer_idx, aux_stream_dict, is_separate_draft_engine=False) - elif model_type in ["qwen3_next", "qwen3_5_text", "qwen3_5_moe_text"]: + 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: From a8234018cb0fd4d481633d329bce5a66c1718ba7 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Fri, 26 Jun 2026 21:20:36 -0700 Subject: [PATCH 08/13] [None][fix] Avoid Python truthiness on 0-D scalar tensors The scalar weight_scale / input_scale lookup used "or" chains: representative = tensors.get("qkv") or next(...) representative_ba = tensors.get("b") or tensors.get("a") For a 0-D torch tensor, bool(tensor) reduces to bool(tensor.item()), so a legitimate zero-valued scalar would silently fall through to the next candidate. Replace with explicit membership checks so the value of the tensor never affects the dispatch. CodeRabbit review nit. Signed-off-by: Mihai Chiorean --- .../_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 c669d7dfc3be..f2966fbc5acf 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 @@ -462,7 +462,7 @@ def _pack_split_projections(self, weights: dict) -> dict: 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.get("qkv") or next( + 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, ( @@ -477,7 +477,7 @@ def _pack_split_projections(self, weights: dict) -> dict: ba_candidates = {"b", "a"} & tensors.keys() if ba_candidates: - representative_ba = tensors.get("b") or tensors.get("a") + 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}" From c7e4248418f0ed1c6420824f3a83516a31d32f19 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 16:16:24 -0700 Subject: [PATCH 09/13] [None][test] Cover Qwen3.6 NVFP4 quant mapper fixes Signed-off-by: Mihai Chiorean --- .../models/test_qwen3_5_moe_weight_mapper.py | 82 +++++++++++++++++++ tests/unittest/llmapi/test_llm_quant.py | 39 +++++++++ 2 files changed, 121 insertions(+) create mode 100644 tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py 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..eff245f1513a --- /dev/null +++ b/tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py @@ -0,0 +1,82 @@ +# 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_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]], + dtype=torch.float8_e4m3fn, + ) + z_fp8 = torch.tensor([[2.0, -1.0]], 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:1].to(torch.float32) * scale, + qkv_fp8[1:2].to(torch.float32) * scale, + qkv_fp8[2:3].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..6f121bb7cba9 100644 --- a/tests/unittest/llmapi/test_llm_quant.py +++ b/tests/unittest/llmapi/test_llm_quant.py @@ -216,6 +216,45 @@ 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, + }, + }, + }, + }, 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"] + assert "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" From db258774047d38068feeb78eed7f8b78c13f6b64 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 16:35:53 -0700 Subject: [PATCH 10/13] [None][style] Apply pre-commit formatting for Qwen3.6 support Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/model_config.py | 3 +- .../checkpoints/hf/qwen3_5_weight_mapper.py | 44 +++++++++---------- .../_torch/models/modeling_qwen3_next.py | 9 ++-- .../_torch/models/modeling_speculative.py | 4 +- .../models/test_qwen3_5_moe_weight_mapper.py | 10 ++--- 5 files changed, 35 insertions(+), 35 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index fc2451d70be5..aeb42a601d85 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -438,7 +438,8 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, _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 + (_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), 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 f2966fbc5acf..7b055b6b0722 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 @@ -39,8 +39,7 @@ # 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], + [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, ) @@ -105,7 +104,7 @@ def _normalize_weight_names(self, weights: dict) -> dict: if key.startswith("model.visual."): continue if key.startswith("model.language_model."): - key = "model." + key[len("model.language_model."):] + key = "model." + key[len("model.language_model.") :] normalized_weights[key] = tensor return normalized_weights @@ -282,10 +281,9 @@ def _dequantize_nvfp4_weight( 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) + 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) @@ -314,8 +312,7 @@ def _dequantize_nvfp4_excluded_weights(self, weights: dict) -> dict: # 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 + 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")] @@ -339,7 +336,6 @@ def _dequantize_nvfp4_excluded_weights(self, weights: dict) -> dict: return updated - def _dequantize_fp8_pertensor_excluded_split_weights(self, weights: dict) -> dict: """Dequantize FP8 per-tensor-scale split linear-attention projections. @@ -379,17 +375,16 @@ def _dequantize_fp8_pertensor_excluded_split_weights(self, weights: dict) -> dic except AttributeError: pass - target_dtype = getattr( - self.config.pretrained_config, "torch_dtype", torch.bfloat16 - ) or torch.bfloat16 + 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 + 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. @@ -406,7 +401,7 @@ def _dequantize_fp8_pertensor_excluded_split_weights(self, weights: dict) -> dic # 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" + scale_key = weight_key[: -len(".weight")] + ".weight_scale" if scale_key not in weights: continue scale = weights[scale_key] @@ -414,13 +409,15 @@ def _dequantize_fp8_pertensor_excluded_split_weights(self, weights: dict) -> dic 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() + (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}" + scale_k = weight_key[: -len(".weight")] + f".{scale_suffix}" updated.pop(scale_k, None) return updated @@ -456,14 +453,15 @@ def _pack_split_projections(self, weights: dict) -> dict: # 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()): + 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"} + 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}, " diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 322e35600d0e..4e2bc6bda341 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -36,13 +36,13 @@ 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 from ..distributed import (AllReduce, AllReduceFusionOp, AllReduceParams, MoEAllReduce, MoEAllReduceParams, allgather) from ..model_config import ModelConfig -from tensorrt_llm.quantization import QuantAlgo from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.fused_moe import (BaseMoeRoutingMethod, MoEWeightLoadingMode, @@ -154,15 +154,16 @@ def __init__( # 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): + 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] + moe_override_quant_config = model_config.quant_config_dict[ + key] break self.experts = create_moe( num_experts=self.num_experts, diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 4867ecf68661..9a5f3a8851e2 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1513,7 +1513,9 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], layer_idx, aux_stream_dict, is_separate_draft_engine=False) - elif model_type in ["qwen3_next", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text"]: + 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/tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py b/tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py index eff245f1513a..0916d37bafec 100644 --- a/tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py +++ b/tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py @@ -19,9 +19,7 @@ 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._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 @@ -57,11 +55,11 @@ def test_fp8_pertensor_linear_attn_weights_are_dequantized_before_pack(): "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), + 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), + "model.layers.0.linear_attn.in_proj_z.input_scale": torch.tensor(1.0, dtype=torch.float32), } packed = mapper.preprocess_weights(weights) From e642f0e4abbce0f80b7575953fb99fb74359c2b3 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 16:39:47 -0700 Subject: [PATCH 11/13] [None][test] Fix Qwen3.6 mapper unit fixture Signed-off-by: Mihai Chiorean --- .../models/test_qwen3_5_moe_weight_mapper.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) 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 index 0916d37bafec..82da0a0e708b 100644 --- a/tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py +++ b/tests/unittest/_torch/models/test_qwen3_5_moe_weight_mapper.py @@ -30,6 +30,7 @@ def _make_mapper() -> Qwen3_5MoeHfWeightMapper: 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, ) @@ -39,7 +40,7 @@ def _make_mapper() -> Qwen3_5MoeHfWeightMapper: quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION), ) mapper = object.__new__(Qwen3_5MoeHfWeightMapper) - mapper.config = model_config + mapper._config = model_config return mapper @@ -47,10 +48,17 @@ 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.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]], 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, @@ -67,9 +75,9 @@ def test_fp8_pertensor_linear_attn_weights_are_dequantized_before_pack(): packed_weight = packed["model.layers.0.linear_attn.in_proj_qkvz.weight"] expected = torch.cat( [ - qkv_fp8[0:1].to(torch.float32) * scale, - qkv_fp8[1:2].to(torch.float32) * scale, - qkv_fp8[2:3].to(torch.float32) * scale, + 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, From 5641db80d633674aece8dc11b1775ed2bc4f370a Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 17:01:34 -0700 Subject: [PATCH 12/13] [None][fix] Exclude prefixed Qwen3.6 lm_head quant configs Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/model_config.py | 16 +++++++++++----- tests/unittest/llmapi/test_llm_quant.py | 7 ++++++- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index aeb42a601d85..f5cf71cdf584 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -445,13 +445,19 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, # 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. - if mixed_quant_configs and "lm_head" in mixed_quant_configs: + 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 = [] - if "lm_head" not in quant_config.exclude_modules: - quant_config.exclude_modules = list( - quant_config.exclude_modules) + ["lm_head"] - del mixed_quant_configs["lm_head"] + 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/tests/unittest/llmapi/test_llm_quant.py b/tests/unittest/llmapi/test_llm_quant.py index 6f121bb7cba9..20dfacb73ab4 100644 --- a/tests/unittest/llmapi/test_llm_quant.py +++ b/tests/unittest/llmapi/test_llm_quant.py @@ -239,6 +239,10 @@ def test_quant_cfg_qwen35_nvfp4_alias_and_prefix_normalization(): "quant_algo": "W4A16_NVFP4", "group_size": 16, }, + "model.language_model.lm_head": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, }, }, }, f) @@ -247,8 +251,9 @@ def test_quant_cfg_qwen35_nvfp4_alias_and_prefix_normalization(): hf_quant_config_file, model_dir, "CUTLASS") assert quant_config.quant_algo == QuantAlgo.MIXED_PRECISION - assert quant_config.exclude_modules == ["lm_head"] + 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 From 0d3b83f6d2bcee42efe985d245b9858506a39abf Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Tue, 30 Jun 2026 21:38:45 -0700 Subject: [PATCH 13/13] [None][docs] Add Qwen3.5 mapper helper docstrings Signed-off-by: Mihai Chiorean --- .../checkpoints/hf/qwen3_5_weight_mapper.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) 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 7b055b6b0722..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 @@ -99,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."): @@ -109,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 @@ -152,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 @@ -190,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, ( @@ -209,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}" @@ -218,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 @@ -230,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) @@ -242,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"): @@ -423,6 +436,12 @@ def _dequantize_fp8_pertensor_excluded_split_weights(self, weights: dict) -> dic 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 @@ -568,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)