[Bugfix] Fix Qwen3-Omni AWQ quantization name mapping (#5687)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/quantization/test_component_routing.py

@@ -9,6 +9,9 @@ import torch
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
 )
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
+    CompressedTensorsConfig,
+)
 from vllm.model_executor.models.utils import WeightsMapper
 
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
@@ -344,7 +347,41 @@ class TestTalkerVisualRouting:
 
 
 # ===================================================================
-# 4. ComponentQuantizationConfig.resolve
+# 4. ComponentQuantizationConfig.apply_vllm_mapper
+# ===================================================================
+
+
+class TestComponentApplyVllmMapper:
+    def test_maps_compressed_tensors_component_and_default(self):
+        thinker = CompressedTensorsConfig(
+            target_scheme_map={},
+            ignore=["thinker.model.layers.0.mlp.gate"],
+            quant_format="pack-quantized",
+        )
+        default = CompressedTensorsConfig(
+            target_scheme_map={},
+            ignore=["talker.model.layers.0.mlp.gate"],
+            quant_format="pack-quantized",
+        )
+        config = ComponentQuantizationConfig(
+            component_configs={"thinker": thinker, "code2wav": None},
+            default_config=default,
+        )
+        mapper = WeightsMapper(
+            orig_to_new_prefix={
+                "thinker.model": "thinker.language_model.model",
+                "talker.model": "talker.language_model.model",
+            }
+        ).get_unstacked_mapper()
+
+        config.apply_vllm_mapper(mapper)
+
+        assert thinker.ignore == ["thinker.language_model.model.layers.0.mlp.gate"]
+        assert default.ignore == ["talker.language_model.model.layers.0.mlp.gate"]
+
+
+# ===================================================================
+# 5. ComponentQuantizationConfig.resolve
 # ===================================================================
 
 

vllm_omni/model_executor/models/qwen3_omni/quantization.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantization mapping helpers for nested Qwen3-Omni stages."""
+
+from vllm.model_executor.models.interfaces import SupportsQuant
+
+_OUTER_MAPPING_APPLIED = "_qwen3_omni_outer_mapping_applied"
+
+
+def apply_outer_quant_config_mapping(model: SupportsQuant) -> None:
+    """Map checkpoint paths once at the unified model boundary."""
+    SupportsQuant._maybe_apply_model_mapping(model)
+    if model.quant_config is not None:
+        setattr(model.quant_config, _OUTER_MAPPING_APPLIED, True)
+
+
+def apply_nested_quant_config_mapping(model: SupportsQuant) -> None:
+    """Avoid applying a stage weight mapper twice to a mapped config."""
+    quant_config = model.quant_config
+    if quant_config is None:
+        return
+    if not getattr(quant_config, _OUTER_MAPPING_APPLIED, False):
+        SupportsQuant._maybe_apply_model_mapping(model)
+        return
+
+    # The outer model maps names, while the nested model still contributes
+    # fused-module metadata needed by quantization scheme matching.
+    if model.packed_modules_mapping is not None:
+        quant_config.packed_modules_mapping.update(model.packed_modules_mapping)
+
+
+class Qwen3OmniNestedSupportsQuant(SupportsQuant):
+    """SupportsQuant variant for models initialized inside the stage wrapper."""
+
+    def _maybe_apply_model_mapping(self) -> None:
+        apply_nested_quant_config_mapping(self)

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -20,12 +20,22 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
 from vllm.config import ModelConfig, VllmConfig
 from vllm.inputs import PromptType, TokensPrompt
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP, SupportsRealtime
+from vllm.model_executor.models.interfaces import (
+    SupportsMRoPE,
+    SupportsMultiModal,
+    SupportsPP,
+    SupportsQuant,
+    SupportsRealtime,
+)
 from vllm.model_executor.models.qwen3_asr_realtime import Qwen3ASRRealtimeBuffer
 from vllm.model_executor.models.qwen3_omni_moe_thinker import (
     Qwen3OmniMoeConditionalGenerationMixin,
 )
-from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
+from vllm.model_executor.models.utils import (
+    WeightsMapper,
+    init_vllm_registered_model,
+    maybe_prefix,
+)
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import MultiModalFeatureSpec
 from vllm.sequence import IntermediateTensors
@@ -39,6 +49,9 @@ from vllm_omni.data_entry_keys import Embeddings, HiddenStates, Ids, OmniPayload
 from vllm_omni.metrics import definitions as defs
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.qwen3_omni.quantization import (
+    apply_outer_quant_config_mapping,
+)
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
     Qwen3OmniMoeThinkerDummyInputsBuilder,
     Qwen3OmniMoeThinkerForConditionalGeneration,
@@ -86,6 +99,7 @@ class Qwen3OmniMoeForConditionalGeneration(
     CustomProcessMixin,
     SupportsMRoPE,
     SupportsRealtime,
+    SupportsQuant,
 ):
     """
     Unified Qwen3 Omni MoE model combining thinker, talker, and code2wav.
@@ -99,6 +113,24 @@ class Qwen3OmniMoeForConditionalGeneration(
         Set `model_stage` in vllm_config to one of: "thinker", "talker", "code2wav"
     """
 
+    # vLLM applies quantization-config name mapping before constructing this
+    # outer stage wrapper.  Expose the final module paths here so checkpoint
+    # ignore lists (for example, the BF16 MoE routers in compressed-tensors
+    # checkpoints) match the nested thinker/talker modules at construction
+    # time.  The keys intentionally omit a trailing dot so exact module names
+    # such as ``thinker.lm_head`` are mapped along with their parameters.
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "thinker.lm_head": "thinker.language_model.lm_head",
+            "thinker.model": "thinker.language_model.model",
+            "talker.model": "talker.language_model.model",
+        }
+    )
+    packed_modules_mapping = Qwen3OmniMoeThinkerForConditionalGeneration.packed_modules_mapping
+
+    def _maybe_apply_model_mapping(self) -> None:
+        apply_outer_quant_config_mapping(self)
+
     realtime_max_tokens = 64
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py

@@ -28,11 +28,14 @@ from vllm.model_executor.models.utils import (  # type: ignore
 )
 
 from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
+from vllm_omni.model_executor.models.qwen3_omni.quantization import (
+    Qwen3OmniNestedSupportsQuant,
+)
 
 logger = init_logger(__name__)
 
 
-class Qwen3OmniMoeCode2Wav(nn.Module):
+class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
     """
     Qwen3 Omni MoE Code2Wav - Converts num_quantizers-layer RVQ codec codes to audio waveform.
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -18,6 +18,9 @@ from vllm.model_executor.models.utils import (
 )
 from vllm.sequence import IntermediateTensors
 
+from vllm_omni.model_executor.models.qwen3_omni.quantization import (
+    Qwen3OmniNestedSupportsQuant,
+)
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp import (
     Qwen3OmniMoeTalkerCodePredictor,
 )
@@ -32,6 +35,7 @@ logger = init_logger(__name__)
 class Qwen3OmniMoeTalkerForConditionalGeneration(
     nn.Module,
     SupportsPP,
+    Qwen3OmniNestedSupportsQuant,
 ):
     """
     Qwen3 Omni MoE Talker - Converts text to audio codec codes.

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -126,6 +126,9 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     _get_video_second_per_grid_t,
     _presampled_videos_hf_kwargs,
 )
+from vllm_omni.model_executor.models.qwen3_omni.quantization import (
+    Qwen3OmniNestedSupportsQuant,
+)
 from vllm_omni.quantization.component_config import (
     PRE_QUANTIZED_METHODS,
     ComponentQuantizationConfig,
@@ -1129,6 +1132,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
     SupportsMRoPE,
     Qwen3OmniMoeConditionalGenerationMixin,
     SupportsTranscription,
+    Qwen3OmniNestedSupportsQuant,
 ):
     # PEFT stores the expert LoRA matrices with an explicit expert dimension.
     is_3d_moe_weight: bool = True

vllm_omni/quantization/component_config.py

@@ -21,6 +21,9 @@ if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import (
         QuantizeMethodBase,
     )
+    from vllm.model_executor.models.utils import (
+        WeightsMapper,
+    )
 
 
 # Pre-quantized checkpoints (modelopt FP8/FP4/MXFP8/mixed) only quantize the
@@ -112,6 +115,14 @@ class ComponentQuantizationConfig(QuantizationConfig):
                 return self._components[comp_prefix]
         return self._default
 
+    def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None:
+        """Apply a weight mapper to every routed quantization config."""
+        for quant_config in self._components.values():
+            if quant_config is not None:
+                quant_config.apply_vllm_mapper(hf_to_vllm_mapper)
+        if self._default is not None:
+            self._default.apply_vllm_mapper(hf_to_vllm_mapper)
+
     def get_name(self) -> str:
         return "component"