[Refactor] Share offline profiler argument parser

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

.buildkite/cuda/test-merge.yml

@@ -3,6 +3,7 @@
 
 env:
   VLLM_WORKER_MULTIPROC_METHOD: spawn
+  VLLM_USE_FLASHINFER_MOE_FP16: 0
   HF_HUB_DOWNLOAD_TIMEOUT: 300
   HF_HUB_ETAG_TIMEOUT: 60
 

.buildkite/cuda/test-nightly.yml

@@ -4,6 +4,7 @@
 
 env:
   VLLM_WORKER_MULTIPROC_METHOD: spawn
+  VLLM_USE_FLASHINFER_MOE_FP16: 0
   HF_HUB_DOWNLOAD_TIMEOUT: 300
   HF_HUB_ETAG_TIMEOUT: 60
 

.buildkite/cuda/test-ready.yml

@@ -3,6 +3,7 @@
 
 env:
   VLLM_WORKER_MULTIPROC_METHOD: spawn
+  VLLM_USE_FLASHINFER_MOE_FP16: 0
   HF_HUB_DOWNLOAD_TIMEOUT: 300
   HF_HUB_ETAG_TIMEOUT: 60
 

.buildkite/cuda/test-weekly.yml

@@ -3,6 +3,7 @@
 
 env:
   VLLM_WORKER_MULTIPROC_METHOD: spawn
+  VLLM_USE_FLASHINFER_MOE_FP16: 0
   HF_HUB_DOWNLOAD_TIMEOUT: 300
   HF_HUB_ETAG_TIMEOUT: 60
 

docs/design/qwen3_omni_tts_performance_optimization.md

@@ -410,8 +410,8 @@ Notes:
 - `stages[].max_num_seqs` controls stage-level batching.
 - Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths.
 - Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior.
-- The Qwen3-Omni deploy config pins `moe_backend: triton` for the Thinker and Talker stages. Triton has been more
- stable and faster than the FlashInfer CUTLASS unquantized MoE backend on recent vLLM rebases.
+- Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster
+ than the FlashInfer CUTLASS unquantized MoE backend on recent vLLM rebases.
 
 #### 2) Enable async chunk
 

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -17,9 +17,9 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 
 The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
 automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
-The bundled Qwen3-Omni setup explicitly sets `moe_backend: triton` for the Thinker and Talker stages. This keeps
-both stages on vLLM's Triton unquantized MoE path and avoids the unstable FlashInfer CUTLASS unquantized MoE path on
-some GPUs.
+The bundled Qwen3-Omni setup defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. This keeps the Thinker & Talker on vLLM's
+Triton unquantized MoE path and avoids the performance regression observed with the FlashInfer CUTLASS unquantized MoE
+backend.
 Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
 
 To explicitly utilize a custom deployment YAML, specify the configuration path:
@@ -87,9 +87,15 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     --stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
 ```
 
-To experiment with another supported MoE backend, set `moe_backend` in the
-deployment YAML for both the Thinker and Talker stages. For the stage-based
-CLI, pass `--moe-backend` directly on each stage command.
+To experiment with the FlashInfer FP16 MoE path, set `VLLM_USE_FLASHINFER_MOE_FP16=1` before launching the server:
+```bash
+VLLM_USE_FLASHINFER_MOE_FP16=1 \
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
+For the stage-based CLI, you usually do **not** need `--stage-overrides` for
+that kind of change. Since each command launches one stage, just pass the knob
+directly on that stage command:
 
 ```bash
 CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -21,6 +21,7 @@ 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
@@ -325,6 +326,7 @@ 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
@@ -378,11 +380,11 @@ def main(args):
         for i, prompt in enumerate(prompts):
             prompt["modalities"] = output_modalities
 
-    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    profiler_enabled = args.profiler_config is not None
     if profiler_enabled and hasattr(omni, "start_profile"):
         omni.start_profile(stages=[0])
     elif profiler_enabled:
-        print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.")
+        print("[Warn] --profiler-config was provided, 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
@@ -458,6 +460,7 @@ 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,6 +14,7 @@ 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"
@@ -45,6 +46,7 @@ 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",
@@ -79,6 +81,7 @@ 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}
@@ -149,8 +152,8 @@ def run_e2e():
 
     sampling_params_list = [gpt_sampling, s2mel_sampling]
 
-    # Start profiling (requires VLLM_TORCH_PROFILER_DIR env var)
-    if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
+    profiler_enabled = args.profiler_config is not None
+    if profiler_enabled:
         print("Starting profiler...")
         omni.start_profile()
 
@@ -158,7 +161,7 @@ def run_e2e():
     outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2]))
 
     # Stop profiling and get results
-    if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
+    if profiler_enabled:
         print("Stopping profiler...")
         profile_results = omni.stop_profile()
         print(f"Profile traces saved to: {profile_results}")

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 24,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 58,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
         EnvironmentVariableCategory.INTERNAL: 2,
     }
@@ -182,7 +182,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
-        ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5,
+        ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 6,
     }
 
 

tests/docs/test_offline_profiler_config.py

@@ -0,0 +1,23 @@
+# 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/engine/test_stage_engine_args.py

@@ -104,6 +104,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
 def _stable_engine_arg_environment(monkeypatch, tmp_path):
     from vllm_omni import platforms
 
+    monkeypatch.delenv("VLLM_USE_FLASHINFER_MOE_FP16", raising=False)
     xcodec_model = tmp_path / "xcodec-model"
     xcodec_model.mkdir()
     monkeypatch.setenv("XCODEC1_PATH", str(xcodec_model))
@@ -136,25 +137,6 @@ def _stable_engine_arg_environment(monkeypatch, tmp_path):
         monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
 
 
-def test_qwen3_omni_defaults_to_triton_moe_backend():
-    engine_args = {"model_arch": "Qwen3OmniMoeForConditionalGeneration"}
-
-    stage_init_utils._maybe_set_qwen3_omni_moe_backend(engine_args)
-
-    assert engine_args["moe_backend"] == "triton"
-
-
-def test_qwen3_omni_preserves_explicit_moe_backend():
-    engine_args = {
-        "model_arch": "Qwen3OmniMoeForConditionalGeneration",
-        "moe_backend": "flashinfer",
-    }
-
-    stage_init_utils._maybe_set_qwen3_omni_moe_backend(engine_args)
-
-    assert engine_args["moe_backend"] == "flashinfer"
-
-
 def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, str]:
     parent_model = tmp_path / "parent-model"
     stage_model = tmp_path / "stage-model"
@@ -535,8 +517,6 @@ def test_typed_engine_args_preserve_explicit_backend_default_overrides(tmp_path)
 
     assert {name: legacy_args[name] for name in expected} == expected
     assert {name: typed_args[name] for name in expected} == expected
-    assert legacy_args["moe_backend"] == "triton"
-    assert typed_args["moe_backend"] == "triton"
 
 
 @pytest.mark.parametrize("stage_scoped", [False, True], ids=["global", "stage-scoped"])

vllm_omni/config/environment_variable_inventory.py

@@ -203,6 +203,7 @@ _MODEL_DEPRECATE_REMOVE = (
     "HIGGS_AUDIO_V2_TOKENIZER_PATH",
     "VLLM_OMNI_LINGBOT_ACTION_ROOT",
     "VLLM_OMNI_VOXCPM_CODE_PATH",
+    "VLLM_USE_FLASHINFER_MOE_FP16",
     "XCODEC1_PATH",
     "model_stage",
 )

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -27,7 +27,6 @@ stages:
     max_num_batched_tokens: 32768
     max_num_seqs: 64
     gpu_memory_utilization: 0.9
-    moe_backend: triton
     trust_remote_code: true
     enable_prefix_caching: false
     devices: "0"
@@ -39,7 +38,6 @@ stages:
     max_num_batched_tokens: 32768
     max_num_seqs: 64
     gpu_memory_utilization: 0.6
-    moe_backend: triton
     trust_remote_code: true
     enable_prefix_caching: false
     devices: "1"

vllm_omni/engine/stage_init_utils.py

@@ -763,14 +763,13 @@ def prepare_engine_environment() -> None:
         pass
 
 
-def _maybe_set_qwen3_omni_moe_backend(engine_args_dict: dict[str, Any]) -> None:
-    """Choose the stable MoE backend when Qwen3-Omni has no explicit choice."""
+def _maybe_set_qwen3_omni_moe_env(engine_args_dict: dict[str, Any]) -> None:
     if (
         engine_args_dict.get("model_arch") == "Qwen3OmniMoeForConditionalGeneration"
-        and engine_args_dict.get("moe_backend", "auto") == "auto"
+        and "VLLM_USE_FLASHINFER_MOE_FP16" not in os.environ
     ):
-        engine_args_dict["moe_backend"] = "triton"
-        logger.info("[stage_init] Set moe_backend=triton for Qwen3-Omni stage")
+        os.environ["VLLM_USE_FLASHINFER_MOE_FP16"] = "0"
+        logger.info("[stage_init] Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni stage")
 
 
 def split_devices_for_replicas(
@@ -1281,8 +1280,9 @@ def _finalize_engine_args_dict(
     engine_args_dict["has_sampling_extra_args"] = has_sampling_extra_args
     engine_args_dict["sampling_extra_args_keys"] = sampling_extra_args_keys
 
-    # Select the typed backend option during stage argument finalization.
-    _maybe_set_qwen3_omni_moe_backend(engine_args_dict)
+    # TODO: Remove this after the performance regression is fixed
+    # Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni to avoid performance regression
+    _maybe_set_qwen3_omni_moe_env(engine_args_dict)
     return engine_args_dict
 
 

vllm_omni/utils/profiler.py

@@ -0,0 +1,37 @@
+# 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"}\'.'
+        ),
+    )