Fix Qwen3-Omni MoE backend selection

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

.buildkite/cuda/test-merge.yml

@@ -3,7 +3,6 @@
 
 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,7 +4,6 @@
 
 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,7 +3,6 @@
 
 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,7 +3,6 @@
 
 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.
-- 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.
+- 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.
 
 #### 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 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.
+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.
 Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
 
 To explicitly utilize a custom deployment YAML, specify the configuration path:
@@ -87,15 +87,9 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     --stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
 ```
 
-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:
+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.
 
 ```bash
 CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \

tests/engine/test_stage_engine_args.py

@@ -104,7 +104,6 @@ _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))
@@ -137,6 +136,25 @@ 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"

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -27,6 +27,7 @@ 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"
@@ -38,6 +39,7 @@ 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,13 +763,14 @@ def prepare_engine_environment() -> None:
         pass
 
 
-def _maybe_set_qwen3_omni_moe_env(engine_args_dict: dict[str, Any]) -> None:
+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."""
     if (
         engine_args_dict.get("model_arch") == "Qwen3OmniMoeForConditionalGeneration"
-        and "VLLM_USE_FLASHINFER_MOE_FP16" not in os.environ
+        and engine_args_dict.get("moe_backend", "auto") == "auto"
     ):
-        os.environ["VLLM_USE_FLASHINFER_MOE_FP16"] = "0"
-        logger.info("[stage_init] Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni stage")
+        engine_args_dict["moe_backend"] = "triton"
+        logger.info("[stage_init] Set moe_backend=triton for Qwen3-Omni stage")
 
 
 def split_devices_for_replicas(
@@ -1280,9 +1281,8 @@ 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
 
-    # 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)
+    # Select the typed backend option during stage argument finalization.
+    _maybe_set_qwen3_omni_moe_backend(engine_args_dict)
     return engine_args_dict