[TRTLLM-10657][fix] Resolve MIXED_PRECISION quant config for DeepSeek W4A8 MoE experts

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -980,6 +980,16 @@ def __init__(self,
             apply_routing=False,
             moe_backend=model_config.moe_backend,
             use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm)
+        # For MIXED_PRECISION, resolve the per-expert quant config (e.g. W4A8_AWQ)
+        # instead of using the ambiguous global MIXED_PRECISION config.
+        # For other cases (e.g. nvfp4, unquantized MTP layers), use
+        # override_quant_config as-is — it already encodes exclusions like MTP.
+        if (override_quant_config is not None and
+                override_quant_config.quant_algo == QuantAlgo.MIXED_PRECISION):
+            expert_quant_config = self._get_experts_quant_config(
+                model_config, layer_idx)
+        else:
+            expert_quant_config = override_quant_config
         self.experts = create_moe(
             num_experts=num_experts,
             routing_method=self.gate.routing_method,
@@ -989,17 +999,13 @@ def __init__(self,
             reduce_results=
             False,  # In both low‑latency and attention‑DP modes, FusedMoE skips the in‑op all‑reduce.
             model_config=model_config,
-            override_quant_config=override_quant_config,
+            override_quant_config=expert_quant_config,
             aux_stream_dict=aux_stream_dict,
             layer_idx=layer_idx,
             # DS-R1 W4A8 is only supported through custom quantization script from
             # examples/quantization/quantize_mixed_precision_moe.py
-            weight_loading_mode=(
-                MoEWeightLoadingMode.W4A8_CUSTOM
-                if self._get_experts_quant_config(
-                    model_config,
-                    layer_idx).layer_quant_mode.is_int4_weight_only_per_group()
-                else MoEWeightLoadingMode.VANILLA),
+            weight_loading_mode=self._expert_weight_loading_mode(
+                expert_quant_config),
         )
 
         self.mapping = model_config.mapping
@@ -1094,6 +1100,19 @@ def _get_experts_quant_config(model_config, layer_idx: int) -> QuantConfig:
         return model_config.quant_config_dict.get(
             f"model.layers.{layer_idx}.mlp.experts", model_config.quant_config)
 
+    @staticmethod
+    def _expert_weight_loading_mode(
+            expert_quant_config: Optional[QuantConfig]) -> MoEWeightLoadingMode:
+        # W4A8_CUSTOM only for the int4-weight-per-group experts produced by
+        # examples/quantization/quantize_mixed_precision_moe.py; everything else
+        # (including the unquantized case, where there is no expert quant config)
+        # uses VANILLA. Guard None: override_quant_config is Optional, so the
+        # resolved expert config can be None on an unquantized layer.
+        if (expert_quant_config is not None and expert_quant_config.
+                layer_quant_mode.is_int4_weight_only_per_group()):
+            return MoEWeightLoadingMode.W4A8_CUSTOM
+        return MoEWeightLoadingMode.VANILLA
+
     @staticmethod
     def _get_shared_experts_quant_config(model_config,
                                          layer_idx: int) -> QuantConfig:
@@ -1300,13 +1319,14 @@ def __init__(self,
             "TRTLLM_DEEPSEEK_EAGER_FUSION_DISABLED", "0") == "0"
         self.enable_fusion &= not self.enable_attention_dp
 
-        # FIXME: incompatible with mixed quantization mode
         quant_config = self._get_decoder_layer_quant_config(
             model_config, layer_idx)
-        self.is_nvfp4 = quant_config.layer_quant_mode.has_nvfp4()
-        assert (
-            quant_config.quant_algo
-            is not QuantAlgo.MIXED_PRECISION), "MIXED_PRECISION is ambiguous"
+        # For MIXED_PRECISION, the global quant_algo doesn't map to a single
+        # QuantMode.  Per-module configs (e.g. expert W4A8_AWQ vs attention
+        # FP8_BLOCK_SCALES) are resolved individually where needed, so we
+        # conservatively set layer-level flags here.
+        self.is_nvfp4 = (quant_config.quant_algo != QuantAlgo.MIXED_PRECISION
+                         and quant_config.layer_quant_mode.has_nvfp4())
 
         self.allreduce = None
         self.moe_allreduce = None

tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py

@@ -463,10 +463,13 @@ def _get_quant_config_dict(self, model_config: ModelConfig) -> Optional[Dict]:
         Extract quantization configuration from model_config
 
         """
-        if model_config.quant_config is None:
+        # Prefer the resolved per-module override (e.g. W4A8_AWQ for experts)
+        # over the global config which may be MIXED_PRECISION.
+        quant_config = getattr(self, "_override_quant_config", None) or model_config.quant_config
+        if quant_config is None:
             return None
 
-        quant_mode = model_config.quant_config.layer_quant_mode
+        quant_mode = quant_config.layer_quant_mode
         return {
             "has_fp8_qdq": quant_mode.has_fp8_qdq()
             if hasattr(quant_mode, "has_fp8_qdq")

tests/unittest/_torch/modeling/test_modeling_deepseekv3.py

@@ -0,0 +1,158 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Per-module quant config resolution for DeepSeek MIXED_PRECISION checkpoints.
+
+DeepSeek-R1-W4AFP8 ships an hf_quant_config with the global
+quant_algo=MIXED_PRECISION, which does not map to a single QuantMode. The MoE
+experts must resolve their own per-module config (W4A8_AWQ) instead of using
+the global one. These tests exercise that resolution on CPU without weights.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from tensorrt_llm._torch.models.modeling_deepseekv3 import Deepseekv3MoE
+from tensorrt_llm._torch.moe.fused_moe import MoEWeightLoadingMode
+from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+pytestmark = pytest.mark.cpu_only
+
+EXPERTS_KEY = "model.layers.{}.mlp.experts"
+
+
+@pytest.fixture
+def mixed_precision_config():
+    return QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION)
+
+
+@pytest.fixture
+def w4a8_awq_config():
+    return QuantConfig(quant_algo=QuantAlgo.W4A8_AWQ, group_size=128)
+
+
+def test_w4a8_awq_config_is_int4_weight_only_per_group(w4a8_awq_config):
+    # This predicate is what selects MoEWeightLoadingMode.W4A8_CUSTOM for the
+    # experts in Deepseekv3MoE, so pin it down explicitly.
+    assert w4a8_awq_config.layer_quant_mode.is_int4_weight_only_per_group()
+
+
+def test_experts_quant_config_resolved_from_per_module_dict(
+    mixed_precision_config, w4a8_awq_config
+):
+    model_config = SimpleNamespace(
+        quant_config=mixed_precision_config,
+        quant_config_dict={EXPERTS_KEY.format(0): w4a8_awq_config},
+    )
+
+    resolved = Deepseekv3MoE._get_experts_quant_config(model_config, 0)
+
+    assert resolved is w4a8_awq_config
+    assert resolved.quant_algo == QuantAlgo.W4A8_AWQ
+    assert resolved.layer_quant_mode.is_int4_weight_only_per_group()
+    assert not mixed_precision_config.layer_quant_mode.is_int4_weight_only_per_group()
+
+
+def test_experts_quant_config_falls_back_to_global_for_unlisted_layer(
+    mixed_precision_config, w4a8_awq_config
+):
+    model_config = SimpleNamespace(
+        quant_config=mixed_precision_config,
+        quant_config_dict={EXPERTS_KEY.format(0): w4a8_awq_config},
+    )
+
+    assert Deepseekv3MoE._get_experts_quant_config(model_config, 1) is mixed_precision_config
+
+
+def test_experts_quant_config_falls_back_to_global_without_dict(mixed_precision_config):
+    model_config = SimpleNamespace(quant_config=mixed_precision_config, quant_config_dict=None)
+
+    assert Deepseekv3MoE._get_experts_quant_config(model_config, 0) is mixed_precision_config
+
+
+def test_expert_weight_loading_mode_w4a8_custom_for_w4a8_awq(w4a8_awq_config):
+    # The resolved W4A8_AWQ expert config selects the custom loading mode; this
+    # is the assignment the MoE construction makes from the resolved config.
+    assert (
+        Deepseekv3MoE._expert_weight_loading_mode(w4a8_awq_config)
+        is MoEWeightLoadingMode.W4A8_CUSTOM
+    )
+
+
+def test_expert_weight_loading_mode_vanilla_for_non_int4(mixed_precision_config):
+    # Neither the ambiguous MIXED_PRECISION global nor a plain FP8 config is
+    # int4-weight-per-group, so both fall to VANILLA.
+    assert (
+        Deepseekv3MoE._expert_weight_loading_mode(mixed_precision_config)
+        is MoEWeightLoadingMode.VANILLA
+    )
+    assert (
+        Deepseekv3MoE._expert_weight_loading_mode(QuantConfig(quant_algo=QuantAlgo.FP8))
+        is MoEWeightLoadingMode.VANILLA
+    )
+
+
+def test_expert_weight_loading_mode_none_is_vanilla():
+    # override_quant_config is Optional, so the resolved expert config is None on
+    # an unquantized layer; the mode selection must not dereference it.
+    assert Deepseekv3MoE._expert_weight_loading_mode(None) is MoEWeightLoadingMode.VANILLA
+
+
+def _bare_configurable_moe(override_quant_config):
+    moe = object.__new__(ConfigurableMoE)
+    moe._override_quant_config = override_quant_config
+    return moe
+
+
+def test_quant_config_dict_prefers_override_over_mixed_precision_global(
+    mixed_precision_config, w4a8_awq_config
+):
+    moe = _bare_configurable_moe(w4a8_awq_config)
+    model_config = SimpleNamespace(quant_config=mixed_precision_config)
+
+    result = ConfigurableMoE._get_quant_config_dict(moe, model_config)
+
+    assert result == {
+        "has_fp8_qdq": False,
+        "has_nvfp4": False,
+        "has_w4afp8": True,
+        "has_fp8_block_scales": False,
+    }
+    # The global MIXED_PRECISION mode would not have flagged w4afp8.
+    assert not mixed_precision_config.layer_quant_mode.is_int4_weight_only_per_group()
+
+
+def test_quant_config_dict_falls_back_to_global_without_override():
+    fp8_config = QuantConfig(quant_algo=QuantAlgo.FP8)
+    moe = _bare_configurable_moe(None)
+    model_config = SimpleNamespace(quant_config=fp8_config)
+
+    result = ConfigurableMoE._get_quant_config_dict(moe, model_config)
+
+    assert result == {
+        "has_fp8_qdq": True,
+        "has_nvfp4": False,
+        "has_w4afp8": False,
+        "has_fp8_block_scales": False,
+    }
+
+
+def test_quant_config_dict_is_none_when_unquantized():
+    moe = _bare_configurable_moe(None)
+    model_config = SimpleNamespace(quant_config=None)
+
+    assert ConfigurableMoE._get_quant_config_dict(moe, model_config) is None