[Docs] Clarify Hunyuan expert remapping scope

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

tests/model_executor/models/test_kv_cache_scale_mapper.py

@@ -67,8 +67,11 @@ class _FakeHunyuanModel(_FakeModel):
             use_cla=False,
         )
 
-    def get_expert_mapping(self) -> tuple[list[object], dict[str, object]]:
+    def get_expert_mapping(self) -> list[object]:
-        return [], {}
+        return []
+
+    def _get_expert_weights_remapping(self) -> dict[str, object]:
+        return {}
 
     def _split_qkv_weight(self, weight: torch.Tensor) -> torch.Tensor:
         return weight
@@ -81,6 +84,49 @@ def test_model_loaders_do_not_call_removed_get_cache_scale(rel_path: str) -> Non
     assert ".get_cache_scale(" not in source
 
 
+@pytest.mark.parametrize(
+    ("module_name", "class_name"),
+    [
+        (
+            "vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3",
+            "HunyuanModel",
+        ),
+        (
+            "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer",
+            "HunyuanImage3Model",
+        ),
+    ],
+)
+def test_hunyuan_expert_mapping_separates_public_and_local_contracts(
+    module_name: str,
+    class_name: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The public expert mapping must be a flat list, not a loader-side pair."""
+    module = __import__(module_name, fromlist=[class_name])
+    expected_mapping = [("param", "weight", 0, "gate_proj")]
+    fake_model = SimpleNamespace(
+        config=SimpleNamespace(num_experts=1),
+        num_redundant_experts=0,
+    )
+    monkeypatch.setattr(module, "_is_moe", lambda config: True)
+    if hasattr(module, "fused_moe_make_expert_params_mapping"):
+        monkeypatch.setattr(module, "fused_moe_make_expert_params_mapping", lambda **kwargs: expected_mapping)
+    else:
+        monkeypatch.setattr(
+            module,
+            "FusedMoE",
+            SimpleNamespace(make_expert_params_mapping=lambda *args, **kwargs: expected_mapping),
+        )
+
+    model_cls = getattr(module, class_name)
+    assert model_cls.get_expert_mapping(fake_model) == expected_mapping
+    assert model_cls._get_expert_weights_remapping(fake_model) == {
+        "gate_proj": ("gate_and_up_proj", 1, 2),
+        "up_proj": ("gate_and_up_proj", 0, 2),
+    }
+
+
 @pytest.mark.parametrize(
     ("module_name", "class_name", "fake_model_cls"),
     [

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -2190,24 +2190,28 @@ class HunyuanImage3Model(nn.Module):
         return torch.concat((q, k, v))
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        if _is_moe(self.config):
+        if not _is_moe(self.config):
-            # Params for weights, fp8 weight scales, fp8 activation scales
+            return []
-            # (param_name, weight_name, expert_id, shard_id)
+
-            fused_moe_expert_mapping = FusedMoE.make_expert_params_mapping(
+        # Params for weights, fp8 weight scales, fp8 activation scales
-                self,
+        # (param_name, weight_name, expert_id, shard_id)
-                ckpt_gate_proj_name="gate_proj",
+        return FusedMoE.make_expert_params_mapping(
-                ckpt_down_proj_name="down_proj",
+            self,
-                ckpt_up_proj_name="up_proj",
+            ckpt_gate_proj_name="gate_proj",
-                num_experts=self.config.num_experts,
+            ckpt_down_proj_name="down_proj",
-                num_redundant_experts=self.num_redundant_experts,
+            ckpt_up_proj_name="up_proj",
-            )
+            num_experts=self.config.num_experts,
-            expert_weights_remapping = {
+            num_redundant_experts=self.num_redundant_experts,
-                "gate_proj": ("gate_and_up_proj", 1, 2),
+        )
-                "up_proj": ("gate_and_up_proj", 0, 2),
+
-            }
+    def _get_expert_weights_remapping(self) -> dict[str, tuple[str, int, int]]:
-            return fused_moe_expert_mapping, expert_weights_remapping
+        """Return the vLLM-Omni-local remapping for Hunyuan checkpoint weights."""
-        else:
+        if not _is_moe(self.config):
-            return [], {}
+            return {}
+        return {
+            "gate_proj": ("gate_and_up_proj", 1, 2),
+            "up_proj": ("gate_and_up_proj", 0, 2),
+        }
 
     # rename for delay load
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
@@ -2243,7 +2247,8 @@ class HunyuanImage3Model(nn.Module):
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
-        expert_params_mapping, expert_weights_remapping = self.get_expert_mapping()
+        expert_params_mapping = self.get_expert_mapping()
+        expert_weights_remapping = self._get_expert_weights_remapping()
 
         # List of unexpected keywords in weight names
         unexpected_keywords = [

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -128,25 +128,29 @@ class HunyuanModel(HunYuanModel):
         v = v.reshape(-1, hidden_size)
         return torch.concat((q, k, v))
 
-    def get_expert_mapping(self) -> tuple[list[tuple[str, str, int, str]], dict[str, tuple[str, int, int]]]:
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        if _is_moe(self.config):
+        if not _is_moe(self.config):
-            # Params for weights, fp8 weight scales, fp8 activation scales
+            return []
-            # (param_name, weight_name, expert_id, shard_id)
+
-            fused_moe_expert_mapping = fused_moe_make_expert_params_mapping(
+        # Params for weights, fp8 weight scales, fp8 activation scales
-                model=self,
+        # (param_name, weight_name, expert_id, shard_id)
-                ckpt_gate_proj_name="gate_proj",
+        return fused_moe_make_expert_params_mapping(
-                ckpt_down_proj_name="down_proj",
+            model=self,
-                ckpt_up_proj_name="up_proj",
+            ckpt_gate_proj_name="gate_proj",
-                num_experts=self.config.num_experts,
+            ckpt_down_proj_name="down_proj",
-                num_redundant_experts=self.num_redundant_experts,
+            ckpt_up_proj_name="up_proj",
-            )
+            num_experts=self.config.num_experts,
-            expert_weights_remapping = {
+            num_redundant_experts=self.num_redundant_experts,
-                "gate_proj": ("gate_and_up_proj", 1, 2),
+        )
-                "up_proj": ("gate_and_up_proj", 0, 2),
+
-            }
+    def _get_expert_weights_remapping(self) -> dict[str, tuple[str, int, int]]:
-            return fused_moe_expert_mapping, expert_weights_remapping
+        """Return the vLLM-Omni-local remapping for Hunyuan checkpoint weights."""
-        else:
+        if not _is_moe(self.config):
-            return [], {}
+            return {}
+        return {
+            "gate_proj": ("gate_and_up_proj", 1, 2),
+            "up_proj": ("gate_and_up_proj", 0, 2),
+        }
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         cla_factor = _get_cla_factor(self.config)
@@ -174,7 +178,8 @@ class HunyuanModel(HunYuanModel):
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
-        expert_params_mapping, expert_weights_remapping = self.get_expert_mapping()
+        expert_params_mapping = self.get_expert_mapping()
+        expert_weights_remapping = self._get_expert_weights_remapping()
 
         # List of unexpected keywords in weight names
         unexpected_keywords = [