[https://nvbugs/5836830][fix] Fix DeepSeekV3 W4A8 mixed precision quantization for MoE experts

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -916,6 +916,16 @@ def __init__(self,
                              fuse_routing_kernel=True,
                              apply_routing=False,
                              moe_backend=model_config.moe_backend)
+        # 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,
@@ -925,17 +935,15 @@ 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=(MoEWeightLoadingMode.W4A8_CUSTOM
+                                 if expert_quant_config.layer_quant_mode.
+                                 is_int4_weight_only_per_group() else
+                                 MoEWeightLoadingMode.VANILLA),
         )
 
         self.mapping = model_config.mapping
@@ -1248,13 +1256,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/modules/fused_moe/configurable_moe.py

@@ -205,6 +205,14 @@ def __init__(
         self.validate_backend(backend)
         self.backend = backend
         self.use_flashinfer = getattr(self.backend, "use_flashinfer", False)
+        # For MIXED_PRECISION, override_quant_config carries the resolved
+        # per-module config (e.g. W4A8_AWQ for MoE experts).  Propagate it
+        # to the backend so that create_weights() picks the correct
+        # quantization strategy instead of the ambiguous global config.
+        self._override_quant_config = override_quant_config
+        if override_quant_config is not None:
+            self.backend.quant_config = override_quant_config
+
         # Sync critical attributes from ConfigurableMoE to backend
         # ConfigurableMoE's super().__init__() was called with real layer_idx and initialized load balancer.
         # Backend was created with init_load_balancer=False and without_comm=True to avoid
@@ -328,10 +336,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/integration/test_lists/test-db/l0_dgx_h200.yml

@@ -42,6 +42,7 @@ l0_dgx_h200:
   - disaggregated/test_disaggregated.py::test_disaggregated_ctxpp4_genpp4[TinyLlama-1.1B-Chat-v1.0]
   - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_ctxtp2ep2pp2_gentp4_one_mtp_block_reuse[DeepSeek-V3-Lite-fp8]
   - unittest/llmapi/test_llm_pytorch.py::test_nemotron_nas_lora
+  - test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_w4afp8_8gpus[DeepSeek-R1-W4AFP8-DeepSeek-R1/DeepSeek-R1-W4AFP8]
 - condition:
     ranges:
       system_gpu_count:

tests/integration/test_lists/waives.txt

@@ -211,7 +211,6 @@ accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_fp8_vswa_reuse SKI
 accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_fp8_guided_decoding_vswa_reuse[xgrammar] SKIP (https://nvbugs/5820497)
 accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_tp4[torch_compile=True] SKIP (https://nvbugs/5821415)
 accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=True] SKIP (https://nvbugs/5821415)
-test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_w4afp8_8gpus[DeepSeek-R1-W4AFP8-DeepSeek-R1/DeepSeek-R1-W4AFP8] SKIP (https://nvbugs/5836830)
 accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664)
 cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mooncake_kvcache-90] SKIP (https://nvbugs/5838199)
 accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-dp4-cutlass-auto] SKIP (https://nvbugs/5838211)