[Hardware][MUSA] Enable Qwen3-Omni ModelOpt FP8 inference (#5671)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/config/test_config_factory.py

@@ -2206,6 +2206,24 @@ class TestBaseConfigInheritance:
 class TestPlatformOverrides:
     """Test platform-specific deploy config overrides."""
 
+    def test_qwen3_omni_musa_clears_parent_quantization_for_audio_stages(self):
+        pipeline = resolve_pipeline_config(
+            "qwen3_omni_moe",
+            Q3_OMNI_ALL_STAGES_HF_CONFIG,
+        )
+        assert isinstance(pipeline, PipelineConfig)
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
+
+        base_stages = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path))
+        assert "hf_overrides" not in base_stages[1].yaml_engine_args
+        assert "hf_overrides" not in base_stages[2].yaml_engine_args
+
+        musa = _apply_platform_overrides(load_deploy_config(deploy_path), platform="musa")
+        musa_stages = merge_pipeline_deploy(pipeline, musa)
+        expected = {"quantization_config": None}
+        assert musa_stages[1].yaml_engine_args["hf_overrides"] == expected
+        assert musa_stages[2].yaml_engine_args["hf_overrides"] == expected
+
     def test_qwen3_tts_rocm_disables_code2wav_outer_cudagraph(self):
         deploy_path = Path(get_deploy_config_path("qwen3_tts.yaml"))
 

tests/worker/test_omni_gpu_model_runner.py

@@ -3,13 +3,88 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.v1.cudagraph_dispatcher import CUDAGraphMode
 
-from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner, _filter_mrope_kwargs_for_model
+from vllm_omni.worker.gpu_model_runner import (
+    OmniGPUModelRunner,
+    _filter_mrope_kwargs_for_model,
+)
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _runner_for_talker_graph_init(
+    *,
+    talker_mtp_graph_safe: bool | None,
+    has_separate_talker: bool = True,
+    model_stage: str = "talker",
+) -> OmniGPUModelRunner:
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = SimpleNamespace(
+        talker=object() if has_separate_talker else None,
+        talker_mtp=object(),
+        model_stage=model_stage,
+    )
+    if talker_mtp_graph_safe is not None:
+        runner.model.talker_mtp_graph_safe = talker_mtp_graph_safe
+    runner.model_config = SimpleNamespace(hf_text_config=SimpleNamespace(hidden_size=4))
+    runner.compilation_config = SimpleNamespace(
+        cudagraph_mode=CUDAGraphMode.FULL,
+        max_cudagraph_capture_size=1,
+    )
+    runner.vllm_config = object()
+    runner.max_num_reqs = 1
+    runner.dtype = torch.float32
+    runner._make_buffer = lambda *args, **kwargs: SimpleNamespace(args=args, kwargs=kwargs)
+    return runner
+
+
+def test_talker_mtp_skips_graph_when_model_declares_unsafe(monkeypatch):
+    runner = _runner_for_talker_graph_init(talker_mtp_graph_safe=False)
+    talker_mtp = runner.model.talker_mtp
+    monkeypatch.setattr(
+        "vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
+        lambda: pytest.fail("graph wrapper must not be selected"),
+    )
+
+    OmniGPUModelRunner._init_talker_mtp(runner)
+
+    assert runner.talker_mtp is talker_mtp
+
+
+@pytest.mark.parametrize("model_stage", ["thinker", "code2wav"])
+def test_non_talker_stage_does_not_use_talker_mtp_graph(monkeypatch, model_stage: str):
+    runner = _runner_for_talker_graph_init(
+        talker_mtp_graph_safe=None,
+        has_separate_talker=False,
+        model_stage=model_stage,
+    )
+    talker_mtp = runner.model.talker_mtp
+    monkeypatch.setattr(
+        "vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
+        lambda: pytest.fail("graph wrapper must not be selected"),
+    )
+
+    OmniGPUModelRunner._init_talker_mtp(runner)
+
+    assert runner.talker_mtp is talker_mtp
+
+
+@pytest.mark.parametrize("talker_mtp_graph_safe", [None, True])
+def test_talker_mtp_uses_graph_for_legacy_or_explicit_safe_model(monkeypatch, talker_mtp_graph_safe):
+    runner = _runner_for_talker_graph_init(talker_mtp_graph_safe=talker_mtp_graph_safe)
+    wrapped = object()
+    monkeypatch.setattr(
+        "vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
+        lambda: lambda *args, **kwargs: wrapped,
+    )
+
+    OmniGPUModelRunner._init_talker_mtp(runner)
+
+    assert runner.talker_mtp is wrapped
+
+
 class DummyBuffer:
     """A minimal buffer wrapper that exposes the `.gpu` attribute."""
 

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -69,6 +69,18 @@ stages:
       repetition_penalty: 1.1
 
 platforms:
+  # The supported Qwen3-Omni ModelOpt checkpoint quantizes the thinker only. Clear the
+  # checkpoint-level metadata for the BF16 audio stages on the verified MUSA
+  # profile before vLLM performs its normal quantization auto-detection.
+  musa:
+    stages:
+      - stage_id: 1
+        hf_overrides:
+          quantization_config: null
+      - stage_id: 2
+        hf_overrides:
+          quantization_config: null
+
   npu:
     stages:
       - stage_id: 0

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -152,6 +152,9 @@ class Qwen3OmniMoeForConditionalGeneration(
                 dtype=torch.long,
             )
         elif self.model_stage == "talker":
+            # The outer wrapper exposes talker_mtp for every stage, but only
+            # the talker stage owns the module that the method invokes.
+            self.talker_mtp_graph_safe = current_omni_platform.supports_talker_mtp_graph_capture()
             multimodal_config.skip_mm_profiling = True
             self.has_preprocess = True
             self.has_postprocess = True

vllm_omni/platforms/interface.py

@@ -155,6 +155,11 @@ class OmniPlatform(Platform):
         """Check if the platform supports torch.compile with inductor backend."""
         raise NotImplementedError
 
+    @classmethod
+    def supports_talker_mtp_graph_capture(cls) -> bool:
+        """Whether a model may capture its dedicated talker MTP graph."""
+        return True
+
     @classmethod
     def has_flash_attn_package(cls) -> bool:
         """Check if a Flash Attention package is available and usable on this platform."""

vllm_omni/platforms/musa/platform.py

@@ -115,6 +115,11 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
         """MUSA supports torch.compile with inductor backend."""
         return True
 
+    @classmethod
+    def supports_talker_mtp_graph_capture(cls) -> bool:
+        """MUSA keeps Qwen3 talker MTP outside its dedicated FULL graph."""
+        return False
+
     @classmethod
     def supports_float64(cls) -> bool:
         """MUSA does not support float64 yet."""

vllm_omni/worker/gpu_model_runner.py

@@ -206,8 +206,12 @@ class OmniGPUModelRunner(GPUModelRunner):
         cudagraph_mode = self.compilation_config.cudagraph_mode
         assert cudagraph_mode is not None
         has_separate_talker = getattr(self.model, "talker", None) is not None
-        talker_mtp_graph_safe = getattr(self.model, "talker_mtp_graph_safe", False)
-        if cudagraph_mode.has_full_cudagraphs() and (has_separate_talker or talker_mtp_graph_safe):
+        talker_mtp_graph_safe = getattr(self.model, "talker_mtp_graph_safe", None)
+        # Preserve the legacy separate-talker default only when the model has
+        # not declared graph safety. An explicit False must disable wrapping.
+        if talker_mtp_graph_safe is None:
+            talker_mtp_graph_safe = has_separate_talker
+        if cudagraph_mode.has_full_cudagraphs() and talker_mtp_graph_safe:
             graph_wrapper_cls = current_omni_platform.get_graph_wrapper_cls()
             self.talker_mtp = graph_wrapper_cls(talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL)
         # TTS exposes mtp_hidden_size; Omni uses hf_text_config.hidden_size.