[Docs] Clarify Hunyuan expert remapping scope

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

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -21,7 +21,6 @@ from vllm.multimodal.media.audio import load_audio
 from vllm.sampling_params import SamplingParams
 
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.profiler import add_profiler_config_arg
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
@@ -326,7 +325,6 @@ def main(args):
     else:
         query_result = query_func()
     args.quantization_config = quantization_config
-    # ``vars(args)`` below forwards profiler_config to Omni.
     omni_kwargs = vars(args).copy()
     # Override CLI --model with the derived model_name.
     omni_kwargs["model"] = model_name
@@ -380,11 +378,11 @@ def main(args):
         for i, prompt in enumerate(prompts):
             prompt["modalities"] = output_modalities
 
-    profiler_enabled = args.profiler_config is not None
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     if profiler_enabled and hasattr(omni, "start_profile"):
         omni.start_profile(stages=[0])
     elif profiler_enabled:
-        print("[Warn] --profiler-config was provided, but current engine does not support profiler controls.")
+        print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.")
     omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
     # Determine output directory: prefer --output-dir; fallback to --output-wav
@@ -460,7 +458,6 @@ def parse_args():
         default=False,
         help="Enable writing detailed statistics (default: disabled)",
     )
-    add_profiler_config_arg(parser)
     parser.add_argument(
         "--stage-init-timeout",
         type=int,

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -14,7 +14,6 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
 from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
-from vllm_omni.utils.profiler import add_profiler_config_arg
 
 # Upstream zero-shot reference clip
 ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
@@ -46,7 +45,6 @@ def run_e2e():
         help="Override the deploy config path. If unset, auto-loads "
         "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
     )
-    add_profiler_config_arg(parser)
     parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
     parser.add_argument(
         "--prompt-text",
@@ -81,7 +79,6 @@ def run_e2e():
         deploy_config=args.deploy_config,
         tokenizer=args.tokenizer,
         log_stats=True,
-        profiler_config=args.profiler_config,
     )
 
     sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
@@ -152,8 +149,8 @@ def run_e2e():
 
     sampling_params_list = [gpt_sampling, s2mel_sampling]
 
-    profiler_enabled = args.profiler_config is not None
-    if profiler_enabled:
+    # Start profiling (requires VLLM_TORCH_PROFILER_DIR env var)
+    if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
         print("Starting profiler...")
         omni.start_profile()
 
@@ -161,7 +158,7 @@ def run_e2e():
     outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2]))
 
     # Stop profiling and get results
-    if profiler_enabled:
+    if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
         print("Stopping profiler...")
         profile_results = omni.stop_profile()
         print(f"Profile traces saved to: {profile_results}")

tests/docs/test_offline_profiler_config.py

@@ -1,23 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from pathlib import Path
-
-import pytest
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-_REPO_ROOT = Path(__file__).resolve().parents[2]
-_EXAMPLES = (
-    "examples/offline_inference/text_to_speech/cosyvoice3/end2end.py",
-    "examples/offline_inference/qwen2_5_omni/end2end.py",
-)
-
-
-@pytest.mark.parametrize("relative_path", _EXAMPLES)
-def test_offline_examples_use_typed_profiler_config(relative_path: str):
-    source = (_REPO_ROOT / relative_path).read_text(encoding="utf-8")
-
-    assert "add_profiler_config_arg(parser)" in source
-    assert "profiler_enabled = args.profiler_config is not None" in source
-    assert "VLLM_TORCH_PROFILER_DIR" not in source

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]]:
-        return [], {}
+    def get_expert_mapping(self) -> list[object]:
+        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):
-            # Params for weights, fp8 weight scales, fp8 activation scales
-            # (param_name, weight_name, expert_id, shard_id)
-            fused_moe_expert_mapping = FusedMoE.make_expert_params_mapping(
-                self,
-                ckpt_gate_proj_name="gate_proj",
-                ckpt_down_proj_name="down_proj",
-                ckpt_up_proj_name="up_proj",
-                num_experts=self.config.num_experts,
-                num_redundant_experts=self.num_redundant_experts,
-            )
-            expert_weights_remapping = {
-                "gate_proj": ("gate_and_up_proj", 1, 2),
-                "up_proj": ("gate_and_up_proj", 0, 2),
-            }
-            return fused_moe_expert_mapping, expert_weights_remapping
-        else:
-            return [], {}
+        if not _is_moe(self.config):
+            return []
+
+        # Params for weights, fp8 weight scales, fp8 activation scales
+        # (param_name, weight_name, expert_id, shard_id)
+        return FusedMoE.make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.num_experts,
+            num_redundant_experts=self.num_redundant_experts,
+        )
+
+    def _get_expert_weights_remapping(self) -> dict[str, tuple[str, int, int]]:
+        """Return the vLLM-Omni-local remapping for Hunyuan checkpoint weights."""
+        if not _is_moe(self.config):
+            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]]]:
-        if _is_moe(self.config):
-            # Params for weights, fp8 weight scales, fp8 activation scales
-            # (param_name, weight_name, expert_id, shard_id)
-            fused_moe_expert_mapping = fused_moe_make_expert_params_mapping(
-                model=self,
-                ckpt_gate_proj_name="gate_proj",
-                ckpt_down_proj_name="down_proj",
-                ckpt_up_proj_name="up_proj",
-                num_experts=self.config.num_experts,
-                num_redundant_experts=self.num_redundant_experts,
-            )
-            expert_weights_remapping = {
-                "gate_proj": ("gate_and_up_proj", 1, 2),
-                "up_proj": ("gate_and_up_proj", 0, 2),
-            }
-            return fused_moe_expert_mapping, expert_weights_remapping
-        else:
-            return [], {}
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        if not _is_moe(self.config):
+            return []
+
+        # Params for weights, fp8 weight scales, fp8 activation scales
+        # (param_name, weight_name, expert_id, shard_id)
+        return fused_moe_make_expert_params_mapping(
+            model=self,
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.num_experts,
+            num_redundant_experts=self.num_redundant_experts,
+        )
+
+    def _get_expert_weights_remapping(self) -> dict[str, tuple[str, int, int]]:
+        """Return the vLLM-Omni-local remapping for Hunyuan checkpoint weights."""
+        if not _is_moe(self.config):
+            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 = [

vllm_omni/utils/profiler.py

@@ -1,37 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from __future__ import annotations
-
-import argparse
-import json
-from typing import Any
-
-
-def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
-    """Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
-    try:
-        config = json.loads(value)
-    except json.JSONDecodeError as e:
-        raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
-    if not isinstance(config, dict):
-        raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
-    return config
-
-
-def parse_profiler_config(value: str) -> dict[str, Any]:
-    """Parse the JSON object passed to ``--profiler-config``."""
-    return parse_json_object(value, flag_name="--profiler-config")
-
-
-def add_profiler_config_arg(parser: argparse.ArgumentParser) -> None:
-    """Add the shared JSON profiler configuration option to ``parser``."""
-    parser.add_argument(
-        "--profiler-config",
-        type=parse_profiler_config,
-        default=None,
-        help=(
-            "JSON profiler config for torch/cuda profiling, e.g. "
-            '\'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.'
-        ),
-    )