[Hardware][Ascend] Accelerate VoxCPM2 LocDiT with NPUGraph (#7275)

Signed-off-by: Big2Wheel <1404207590@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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)])

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

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

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():

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]

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)

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