From e04210d690d619d6da34cbce3ae98e6fa31398bb Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Tue, 11 Aug 2026 17:19:09 +0800 Subject: [PATCH] [Hardware][MUSA] Enable Qwen3-Omni ModelOpt FP8 inference (#5671) Signed-off-by: Xiaodong Ye Co-authored-by: Hongsheng Liu Co-authored-by: amy-why-3459 --- tests/config/test_config_factory.py | 18 +++++ tests/worker/test_omni_gpu_model_runner.py | 77 ++++++++++++++++++- vllm_omni/deploy/qwen3_omni_moe.yaml | 12 +++ .../models/qwen3_omni/qwen3_omni.py | 3 + vllm_omni/platforms/interface.py | 5 ++ vllm_omni/platforms/musa/platform.py | 5 ++ vllm_omni/worker/gpu_model_runner.py | 8 +- 7 files changed, 125 insertions(+), 3 deletions(-) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index e760548f8..eb774bb18 100644 --- a/tests/config/test_config_factory.py +++ b/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")) diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index 701e3b589..949a2c7a8 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/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.""" diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index 5bcbb0fe3..76878804d 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/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 diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index b65cfdad0..f5a1551a7 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/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 diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index fdd31ad21..7759cb6d1 100644 --- a/vllm_omni/platforms/interface.py +++ b/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.""" diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 02ce155ee..df52905c9 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/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.""" diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 9fb781dc4..0e0d6a46e 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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.