[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/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."""