[#14287][feat] Add Qwen3.5/3.6 MoE NVFP4 + MTP support for SM120/SM121

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:

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(

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:

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
 
 

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:

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

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"