From 7ce935df2dca0a69b7fb4c9ce72c9734d7f7f8a8 Mon Sep 17 00:00:00 2001
From: wangqian <601731555@qq.com>
Date: Tue, 14 Jul 2026 13:19:13 +0800
Subject: [PATCH] [Bugfix][Gemma4] Fix ModelOpt mixed-precision MoE mapping

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: wangqian <601731555@qq.com>
---
 tests/quantization/test_modelopt.py     | 32 +++++++++++++++++++++++++
 vllm/model_executor/models/gemma4.py    |  8 ++++++-
 vllm/model_executor/models/gemma4_mm.py |  9 ++++---
 3 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py
index a08e14c53c8c..d2541dcc9107 100644
--- a/tests/quantization/test_modelopt.py
+++ b/tests/quantization/test_modelopt.py
@@ -165,6 +165,38 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
     assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None
 
 
+def test_modelopt_mixed_precision_composes_gemma4_mappers():
+    from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
+    from vllm.model_executor.models.gemma4_mm import (
+        Gemma4ForConditionalGeneration,
+    )
+
+    config = _mixed_precision_config(
+        {
+            "model.language_model.layers.0.experts": {
+                "quant_algo": "NVFP4",
+                "group_size": 16,
+            },
+            "model.language_model.layers.1.moe.experts.gate_up_proj": {
+                "quant_algo": "NVFP4",
+                "group_size": 16,
+            },
+        }
+    )
+
+    config.apply_vllm_mapper(
+        Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
+    )
+    config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())
+
+    expected_prefix = "language_model.model.layers.0.moe.experts"
+    assert set(config.quantized_layers) == {
+        expected_prefix,
+        "language_model.model.layers.1.moe.gate_up_proj",
+    }
+    assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
+
+
 def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
     from vllm.model_executor.layers.linear import LinearBase
 
diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py
index 63d97fbd1b50..e0c3130bddaa 100644
--- a/vllm/model_executor/models/gemma4.py
+++ b/vllm/model_executor/models/gemma4.py
@@ -84,6 +84,12 @@
 
 logger = init_logger(__name__)
 
+_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper(
+    orig_to_new_regex={
+        re.compile(r"(?<!\.moe)\.experts$"): ".moe.experts",
+    }
+)
+
 
 def _remap_gemma4_expert_weight_name(name: str) -> str:
     return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
@@ -1508,7 +1514,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 class Gemma4ForCausalLM(
     nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
 ):
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
         orig_to_new_prefix={
             # Gemma4ForConditionalGeneration already loads the text stack
             # from `model.language_model.*`. We reuse that same checkpoint
diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py
index d733a181a004..b0169cbc6ddf 100644
--- a/vllm/model_executor/models/gemma4_mm.py
+++ b/vllm/model_executor/models/gemma4_mm.py
@@ -40,7 +40,10 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import ReplicatedLinear
-from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
+from vllm.model_executor.models.gemma4 import (
+    _GEMMA4_EXPERT_PARENT_MAPPER,
+    Gemma4ForCausalLM,
+)
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.transformers.utils import recursive_replace_linear
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -998,7 +1001,7 @@ class Gemma4ForConditionalGeneration(
     }
 
     # Maps checkpoint prefixes to vLLM module paths.
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
         orig_to_new_prefix={
             # vision tower
             "model.vision_tower": "vision_tower",
@@ -1010,7 +1013,7 @@ class Gemma4ForConditionalGeneration(
             "model.language_model.": "language_model.model.",
             "lm_head.": "language_model.lm_head.",
             "model": "language_model.model",
-        }
+        },
     )
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):