From 9097fa654742a70801a976d8f140d337a62c1266 Mon Sep 17 00:00:00 2001 From: yulun <100981785+Big2Wheel@users.noreply.github.com> Date: Wed, 30 Sep 2026 10:11:22 +0800 Subject: [PATCH] [Hardware][Ascend] Accelerate VoxCPM2 LocDiT with NPUGraph (#7275) Signed-off-by: Big2Wheel <1404207590@qq.com> Co-authored-by: Hongsheng Liu --- ...ni_ar_scheduler_unified_decode_deferral.py | 4 +- ...st_runner_assisted_unified_graph_static.py | 3 +- .../npu/test_voxcpm2_loc_dit_npugraph.py | 158 ++++++++++++++++++ .../models/voxcpm2/scheduler.py | 6 +- .../models/voxcpm2/voxcpm2_talker.py | 25 +-- .../platforms/npu/models/voxcpm2_talker.py | 94 +++++++++++ vllm_omni/platforms/npu/platform.py | 4 + 7 files changed, 277 insertions(+), 17 deletions(-) create mode 100644 tests/platforms/npu/test_voxcpm2_loc_dit_npugraph.py create mode 100644 vllm_omni/platforms/npu/models/voxcpm2_talker.py diff --git a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py index de68f7d5a..b55a668dd 100644 --- a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py +++ b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py @@ -57,7 +57,7 @@ class _MockRequest: @pytest.fixture(autouse=True) def _mock_cuda_graph_platform(monkeypatch) -> None: - monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: True) + monkeypatch.setattr(voxcpm2_scheduler_mod.omni_platform.current_omni_platform, "is_cuda", lambda: True) def _make_scheduler( @@ -134,7 +134,7 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_with_deterministic_noise() def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypatch) -> None: - monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: False) + monkeypatch.setattr(voxcpm2_scheduler_mod.omni_platform.current_omni_platform, "is_cuda", lambda: False) scheduler = _make_scheduler() scheduler.running = [_MockRequest("decode")] scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]) diff --git a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py index 0c2d9ee67..03c00f122 100644 --- a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py +++ b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py @@ -192,9 +192,10 @@ def test_voxcpm2_scheduler_policy_stays_model_local(): assert "def _should_defer_waiting_admission(self) -> bool:" in common_source assert "class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler)" in voxcpm2_scheduler_source + assert "import vllm_omni.platforms as omni_platform" in voxcpm2_scheduler_source assert "from .runtime_config import _VoxCPM2RuntimeConfig" in voxcpm2_scheduler_source assert "_VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)" in voxcpm2_scheduler_source - assert "unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())" in voxcpm2_scheduler_source + assert "use_cuda_graph=omni_platform.current_omni_platform.is_cuda()" in voxcpm2_scheduler_source assert "_should_defer_waiting_for_unified_decode_graph" in voxcpm2_scheduler_source assert "def schedule(" not in voxcpm2_scheduler_source assert "create_request_queue" not in voxcpm2_scheduler_source diff --git a/tests/platforms/npu/test_voxcpm2_loc_dit_npugraph.py b/tests/platforms/npu/test_voxcpm2_loc_dit_npugraph.py new file mode 100644 index 000000000..d34a73209 --- /dev/null +++ b/tests/platforms/npu/test_voxcpm2_loc_dit_npugraph.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the Ascend VoxCPM2 LocDiT NPUGraph adapter.""" + +import sys +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.platforms.npu.models import voxcpm2_talker as npu_adapter + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeEstimator(torch.nn.Module): + def forward(self, x, mu, t, cond, dt): + return x + mu + t + cond + dt + + +class FakeTalker: + def __init__(self, *, vllm_config=None, prefix="") -> None: + self.vllm_config = vllm_config + self.prefix = prefix + self.estimator = FakeEstimator() + self.tts = SimpleNamespace(feat_decoder=SimpleNamespace(estimator=self.estimator)) + + +class FakeTalkerSubclass(FakeTalker): + pass + + +class FakeGraphRunner: + instances = [] + supported = True + + def __init__(self, *, max_graphs, component_name, disable_config_hint) -> None: + self.max_graphs = max_graphs + self.component_name = component_name + self.disable_config_hint = disable_config_hint + self.calls = [] + self.__class__.instances.append(self) + + @classmethod + def is_supported(cls) -> bool: + return cls.supported + + def run(self, operation, inputs, constants, compute): + self.calls.append((operation, inputs, constants)) + return compute(*inputs) + + +def test_model_patch_wraps_init_and_is_idempotent(monkeypatch) -> None: + calls = [] + monkeypatch.setattr(npu_adapter, "setup_voxcpm2_loc_dit_npu_graph", calls.append) + monkeypatch.setattr(npu_adapter, "_PATCHED", False) + monkeypatch.setattr(npu_adapter, "_original_init", None) + + original_init = FakeTalker.__init__ + fake_module = SimpleNamespace(VoxCPM2TalkerForConditionalGeneration=FakeTalker) + monkeypatch.setitem( + sys.modules, + "vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker", + fake_module, + ) + + try: + npu_adapter.apply_voxcpm2_talker_patch() + npu_adapter.apply_voxcpm2_talker_patch() + model = FakeTalker(vllm_config="config", prefix="talker") + finally: + FakeTalker.__init__ = original_init + + assert calls == [model] + assert model.vllm_config == "config" + assert model.prefix == "talker" + + +def test_npu_platform_registers_patch_and_keeps_worker_cls(monkeypatch) -> None: + pytest.importorskip( + "vllm_ascend", + reason="NPU platform registration requires the optional vllm-ascend plugin", + ) + from vllm_ascend import utils as ascend_utils + + from vllm_omni.platforms.npu import _310p + from vllm_omni.platforms.npu.models import ( + minicpmo_4_5_code2wav, + qwen3_tts, + qwen3_tts_tokenizer_v2, + ) + from vllm_omni.platforms.npu.platform import NPUOmniPlatform + + calls = [] + monkeypatch.setattr(ascend_utils, "adapt_patch", lambda **kwargs: None) + monkeypatch.setattr(_310p, "apply_patches", lambda: None) + monkeypatch.setattr( + minicpmo_4_5_code2wav, + "apply_minicpmo_4_5_code2wav_patch", + lambda: None, + ) + monkeypatch.setattr(qwen3_tts, "apply_qwen3_tts_patches", lambda: None) + monkeypatch.setattr( + qwen3_tts_tokenizer_v2, + "apply_qwen3_tts_tokenizer_v2_patch", + lambda: None, + ) + monkeypatch.setattr( + npu_adapter, + "apply_voxcpm2_talker_patch", + lambda: calls.append("voxcpm2"), + ) + + NPUOmniPlatform() + worker_cls = NPUOmniPlatform.get_omni_ar_worker_cls() + + assert worker_cls == "vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker" + assert calls == ["voxcpm2"] + + +def test_loc_dit_npugraph_supports_wrapped_subclass_and_wraps_once(monkeypatch) -> None: + FakeGraphRunner.instances.clear() + FakeGraphRunner.supported = True + monkeypatch.setattr( + npu_adapter, + "_get_npu_exact_graph_runner_cls", + lambda: FakeGraphRunner, + ) + model = FakeTalkerSubclass() + wrapped_model = SimpleNamespace(module=model) + + npu_adapter.setup_voxcpm2_loc_dit_npu_graph(wrapped_model) + npu_adapter.setup_voxcpm2_loc_dit_npu_graph(wrapped_model) + + inputs = tuple(torch.full((2, 3), value) for value in range(5)) + output = model.estimator(*inputs) + runner = FakeGraphRunner.instances[0] + assert torch.equal(output, sum(inputs[1:], start=inputs[0])) + assert runner.max_graphs == 8 + assert runner.component_name == "VoxCPM2 LocDiT" + assert runner.calls == [("forward", inputs, ())] + assert len(FakeGraphRunner.instances) == 1 + + +def test_loc_dit_npugraph_falls_back_when_apis_are_unavailable(monkeypatch) -> None: + FakeGraphRunner.instances.clear() + FakeGraphRunner.supported = False + monkeypatch.setattr( + npu_adapter, + "_get_npu_exact_graph_runner_cls", + lambda: FakeGraphRunner, + ) + model = FakeTalker() + + npu_adapter.setup_voxcpm2_loc_dit_npu_graph(model) + + assert not hasattr(model.estimator, "_voxcpm2_npu_graph_runner") + FakeGraphRunner.supported = True diff --git a/vllm_omni/model_executor/models/voxcpm2/scheduler.py b/vllm_omni/model_executor/models/voxcpm2/scheduler.py index a780c6955..05d73653f 100644 --- a/vllm_omni/model_executor/models/voxcpm2/scheduler.py +++ b/vllm_omni/model_executor/models/voxcpm2/scheduler.py @@ -5,8 +5,8 @@ from __future__ import annotations from vllm.v1.request import RequestStatus +import vllm_omni.platforms as omni_platform from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler -from vllm_omni.platforms import current_omni_platform from .runtime_config import _VoxCPM2RuntimeConfig @@ -23,7 +23,9 @@ class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler): def _unified_decode_graph_enabled(self) -> bool: runtime_config = _VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config) - return runtime_config.unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda()) + return runtime_config.unified_decode_graph_available( + use_cuda_graph=omni_platform.current_omni_platform.is_cuda() + ) def _should_defer_waiting_for_unified_decode_graph(self) -> bool: if not self._unified_decode_graph_enabled(): diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 47c6f1dd3..4894ac5f7 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -36,8 +36,8 @@ from vllm.multimodal.audio import AudioResampler from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +import vllm_omni.platforms as omni_platform from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.platforms import current_omni_platform from vllm_omni.utils.speaker_cache import ( get_speaker_cache, iter_custom_voice_profiles, @@ -377,7 +377,7 @@ class _PerfTimer: self._enabled = enabled self._device_module = None if enabled: - device_module = torch.get_device_module(current_omni_platform.get_torch_device()) + device_module = torch.get_device_module(omni_platform.current_omni_platform.get_torch_device()) if not hasattr(device_module, "Event"): logger.warning_once( "VoxCPM2 profiler disabled: the current device does not provide accelerator timing events" @@ -904,7 +904,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): # the checkpoint download/read cost at construction time; DummyModelLoader # will then randomize the just-loaded _tts params — this is intended. model_path = vllm_config.model_config.model - self._device = current_omni_platform.get_torch_device() + self._device = omni_platform.current_omni_platform.get_torch_device() VoxCPM = import_voxcpm2_core() native = VoxCPM.from_pretrained(model_path, load_denoiser=False, optimize=False) self._tts: nn.Module = native.tts_model.to(self._device) @@ -929,8 +929,8 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): # for sliding-window streaming decode. 12 matches the nanovllm reference # implementation and covers the longest VAE decoder receptive field. self._n_decode_pad_frames = 12 - use_cuda_graph = current_omni_platform.is_cuda() - self._enable_torch_compile = current_omni_platform.supports_torch_inductor() + use_cuda_graph = omni_platform.current_omni_platform.is_cuda() + self._enable_torch_compile = omni_platform.current_omni_platform.supports_torch_inductor() self._compile_vae = self._enable_torch_compile self._max_decode_steps = 2000 self._max_batch_size = getattr(vllm_config.scheduler_config, "max_num_seqs", 4) @@ -1579,7 +1579,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): return sr_cond def _run_vae_decode(self, feat: torch.Tensor) -> torch.Tensor: - if feat.device.type != current_omni_platform.device_type: + if feat.device.type != omni_platform.current_omni_platform.device_type: return self.tts.audio_vae.decode(feat) sr_cond = self._get_vae_decode_sr_cond(feat.device) @@ -2355,7 +2355,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): def _run_cfm(self, dit_h: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: with _NvtxRange("voxcpm2.cfm"): if self._cfm_buffers is not None: - if self._enable_cfm_cuda_graph and dit_h.device.type == current_omni_platform.device_type: + if self._enable_cfm_cuda_graph and dit_h.device.type == omni_platform.current_omni_platform.device_type: return self._run_cfm_cuda_graph(dit_h, cond).transpose(1, 2) return _optimized_solve_euler( self.tts.feat_decoder, @@ -2379,7 +2379,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): def _run_cfm_for_state(self, state: _RequestState, dit_h: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: with _NvtxRange("voxcpm2.cfm"): if self._cfm_buffers is not None: - if self._enable_cfm_cuda_graph and dit_h.device.type == current_omni_platform.device_type: + if self._enable_cfm_cuda_graph and dit_h.device.type == omni_platform.current_omni_platform.device_type: if self._has_deterministic_cfm_noise(state) and not self._enable_cfm_prealloc_output: graph = self._get_cfm_cuda_graph(dit_h, cond) with _NvtxRange("voxcpm2.cfm.graph_copy_mu"): @@ -2715,7 +2715,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): def _enqueue_delayed_audio_copy(self, state: _RequestState, audio: torch.Tensor) -> None: src = audio.detach().contiguous() - if src.device.type != current_omni_platform.device_type: + if src.device.type != omni_platform.current_omni_platform.device_type: state.pending_audio_copies.append(_PendingAudioCopy(host=src.cpu().contiguous())) return @@ -2803,7 +2803,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if state.is_stopping or state.precomputed_is_stopping is not None: continue stop_logits = state.precomputed_stop_logits - if stop_logits is None or stop_logits.device.type != current_omni_platform.device_type: + if stop_logits is None or stop_logits.device.type != omni_platform.current_omni_platform.device_type: continue pending.append((state, stop_logits)) if not pending: @@ -3085,7 +3085,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): ready_req_ids = list(audio_by_req) chunks = [audio_by_req[req_id].reshape(-1) for req_id in ready_req_ids] if self._coalesce_audio_d2h and any( - chunk.device.type == current_omni_platform.device_type for chunk in chunks + chunk.device.type == omni_platform.current_omni_platform.device_type for chunk in chunks ): sizes = [int(chunk.numel()) for chunk in chunks] merged = torch.cat(chunks, dim=0) if len(chunks) > 1 else chunks[0] @@ -3096,7 +3096,8 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): mm["sr"] = [sr for _ in ready_req_ids] mm["meta"] = {"req_id": ready_req_ids, "sparse_audio": ["1"]} elif self._coalesce_audio_d2h and any( - audio.device.type == current_omni_platform.device_type for audio in audio_by_req.values() + audio.device.type == omni_platform.current_omni_platform.device_type + for audio in audio_by_req.values() ): ready_req_ids = list(audio_by_req) chunks = [audio_by_req[req_id].reshape(-1) for req_id in ready_req_ids] diff --git a/vllm_omni/platforms/npu/models/voxcpm2_talker.py b/vllm_omni/platforms/npu/models/voxcpm2_talker.py new file mode 100644 index 000000000..5ef3b2aec --- /dev/null +++ b/vllm_omni/platforms/npu/models/voxcpm2_talker.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Enable VoxCPM2 LocDiT NPUGraph acceleration on Ascend.""" + +from __future__ import annotations + +from types import MethodType +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_MAX_GRAPHS = 8 +_PATCHED = False +_original_init = None + + +def _patched_init(self, *, vllm_config: Any, prefix: str = "") -> None: + assert _original_init is not None + _original_init(self, vllm_config=vllm_config, prefix=prefix) + setup_voxcpm2_loc_dit_npu_graph(self) + + +def _get_npu_exact_graph_runner_cls(): + # Import graph tooling only in the engine worker. Importing it in vLLM's + # short-lived model-inspection subprocess initializes the NPU runtime and + # can make that subprocess abort during teardown. + from vllm_omni.platforms.npu.graph_tools import NPUExactGraphRunner + + return NPUExactGraphRunner + + +def apply_voxcpm2_talker_patch() -> None: + """Install the Ascend LocDiT adapter before VoxCPM2 is constructed.""" + global _PATCHED, _original_init + if _PATCHED: + return + + # The model defers resolving ``current_omni_platform`` until construction, + # so importing it while NPUOmniPlatform is initialized is cycle-free. + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + VoxCPM2TalkerForConditionalGeneration, + ) + + _original_init = VoxCPM2TalkerForConditionalGeneration.__init__ + VoxCPM2TalkerForConditionalGeneration.__init__ = _patched_init # type: ignore[method-assign] + _PATCHED = True + logger.debug("Applied NPU patch for VoxCPM2 LocDiT") + + +def setup_voxcpm2_loc_dit_npu_graph(model: object) -> None: + """Wrap VoxCPM2's tensor-only LocDiT estimator with NPUGraph.""" + wrapped_model = getattr(model, "module", None) + talker = wrapped_model if wrapped_model is not None else model + tts = getattr(talker, "tts", None) + feat_decoder = getattr(tts, "feat_decoder", None) + estimator = getattr(feat_decoder, "estimator", None) + if estimator is None: + raise TypeError("expected a VoxCPM2 talker with tts.feat_decoder.estimator") + if getattr(estimator, "_voxcpm2_npu_graph_runner", None) is not None: + return + + graph_runner = _get_npu_exact_graph_runner_cls()( + max_graphs=_MAX_GRAPHS, + component_name="VoxCPM2 LocDiT", + disable_config_hint="disable the VoxCPM2 NPU model patch", + ) + if not graph_runner.is_supported(): + logger.warning("VoxCPM2 LocDiT NPUGraph APIs are unavailable; using eager execution") + return + + original_forward = estimator.forward + + def _forward( + module, + x: torch.Tensor, + mu: torch.Tensor, + t: torch.Tensor, + cond: torch.Tensor, + dt: torch.Tensor, + ) -> torch.Tensor: + del module + return graph_runner.run( + "forward", + (x, mu, t, cond, dt), + (), + lambda *inputs: (original_forward(*inputs),), + )[0] + + estimator.forward = MethodType(_forward, estimator) + estimator._voxcpm2_npu_graph_runner = graph_runner + logger.info("VoxCPM2 LocDiT NPUGraph replay enabled (max_graphs=%d)", _MAX_GRAPHS) diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 24e29f342..8165c1a42 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -68,11 +68,15 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): from vllm_omni.platforms.npu.models.qwen3_tts_tokenizer_v2 import ( apply_qwen3_tts_tokenizer_v2_patch, ) + from vllm_omni.platforms.npu.models.voxcpm2_talker import ( + apply_voxcpm2_talker_patch, + ) adapt_patch(is_global_patch=True) apply_minicpmo_4_5_code2wav_patch() apply_qwen3_tts_patches() apply_qwen3_tts_tokenizer_v2_patch() + apply_voxcpm2_talker_patch() apply_310p_patches() @classmethod