[MiniCPM-o][NPU] Enable MiniCPM-o Code2Wav NPUGraph replay (#5604)

Signed-off-by: Zhou248 <62542847+Zhou248@users.noreply.github.com>

.buildkite/npu/test-npu-ready.yml

@@ -1,4 +1,15 @@
 steps:
+  - label: "Platforms Test"
+    key: ready-platforms-test
+    depends_on: upload-ready-pipeline
+    mirror_hardwares: a3_npu_2
+    commands:
+      - >-
+        pytest -s -v
+        tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
+        -m 'core_model and npu and A3 and omni'
+        --run-level 'core_model'
+
   - group: ":card_index_dividers: Omni Model Test"
     key: ready-omni-group
     depends_on: upload-ready-pipeline

recipes/OpenBMB/MiniCPM-o-4_5.md

@@ -71,6 +71,35 @@ single-GPU entry point; `minicpmo_4_5_2gpu.yaml` is the recommended
 two-GPU profile. The removed fused two-stage implementation is not retained as
 a fallback because it would duplicate state machines and correctness paths.
 
+### Graph execution
+
+Graph boundaries follow each stage's state model:
+
+| Stage | Graph mode on Ascend | Eager boundary |
+| --- | --- | --- |
+| 0 Thinker | vLLM `PIECEWISE` | multimodal preprocessing and output routing |
+| 1 Talker | vLLM `PIECEWISE` | conditioning preprocessing, codec sampling, request-state updates |
+| 2 Code2Wav | inner exact-shape NPUGraph for the CFM DiT estimator | encoder, timestep embedding, HiFT/RNG, request parsing, stream-state commit |
+
+Stage 2 keeps `enforce_eager: true` for the generation runner because its
+Python request metadata and per-request cache dictionaries are not valid
+outer-graph inputs. The backend captures only the deterministic CFM DiT
+estimator, with an exact graph key for each tensor shape. Timestep embedding
+stays eager because the upstream implementation creates a CPU frequency tensor;
+HiFT stays eager because it generates random phase. The graph cache stores up
+to 32 entries by default, after which unseen shapes run eagerly.
+
+An ACL graph capture failure is fatal to that Stage-2 process because older
+torch-npu releases can leave allocator and RNG capture state invalid. Restart
+the service after a capture failure. To run without capture, set
+`code2wav_enable_npu_graph: false` under the Stage-2
+`platforms.npu.stages[].additional_config` block before startup. Tune the cache
+limit there with `code2wav_max_npu_graphs`. Graph mode also requires
+`ASCEND_LAUNCH_BLOCKING` to be unset or set to `0`.
+
+The inner NPUGraph is independent of the outer runner setting, so do not use a
+global `--enforce-eager` override when Stage 0/1 `PIECEWISE` replay is desired.
+
 ## GPU
 
 ### 1 x GPU (default — single command)
@@ -211,9 +240,9 @@ speech output (TTS)"** checkbox on / off.
   model processes still share one CUDA device.
 - `--trust-remote-code` is required — the HF repo ships a custom
   `MiniCPMO` config / model class.
-- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains
-  eager because its request-owned Flow/HiFT caches and variable chunk/cache
-  shapes are not yet exposed through a static exact-shape graph wrapper.
+- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 keeps its
+  request orchestration eager while using inner CFM and HiFT CUDA Graphs. On
+  Ascend, the exact-shape CFM DiT estimator instead uses inner NPUGraph replay.
 - All default stages use `max_num_seqs: 4` to reduce cross-process GPU
   contention. Talker AR
   state and Code2Wav caches are request-owned; Code2Wav batches only
@@ -224,8 +253,8 @@ speech output (TTS)"** checkbox on / off.
 - `StageRequestStats.batch_size` is a request-scoped placeholder, not the
   scheduler's execution batch.
 - Single-GPU co-location trades throughput for hardware density: Stage 0/1
-  CUDA Graph replay and eager Stage 2 vocoder kernels compete across three
-  CUDA contexts. Use the 8x4090 config or a custom multi-GPU mapping for
+  CUDA Graph replay and Stage 2 vocoder kernels compete across three CUDA
+  contexts. Use the 8x4090 config or a custom multi-GPU mapping for
   throughput-sensitive serving.
 
 ### 8 x RTX 4090 24GB (consumer-GPU layout)

tests/model_executor/models/minicpmo_4_5/test_pipeline.py

@@ -24,6 +24,7 @@ from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import (
     PipelineConfig,
     StageExecutionType,
+    _apply_platform_overrides,
     load_deploy_config,
     merge_pipeline_deploy,
 )
@@ -184,6 +185,30 @@ class TestDeployTopology:
                 0.35,
             ]
 
+    @pytest.mark.parametrize(
+        "filename",
+        [
+            "minicpmo_4_5.yaml",
+            "minicpmo_4_5_2gpu.yaml",
+            "minicpmo_4_5_3gpu.yaml",
+        ],
+    )
+    def test_npu_graph_configuration_is_stage_scoped(self, filename: str) -> None:
+        deploy = load_deploy_config(_DEPLOY_DIR / filename)
+        connector_extra = deploy.connectors["connector_of_shared_memory"]["extra"]
+        assert "code2wav_enable_npu_graph" not in connector_extra
+        assert "code2wav_max_npu_graphs" not in connector_extra
+
+        deploy = _apply_platform_overrides(deploy, platform="npu")
+        stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy)
+        assert stages[0].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE"
+        assert stages[1].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE"
+        assert stages[2].yaml_engine_args["enforce_eager"] is True
+        assert stages[2].yaml_engine_args["additional_config"] == {
+            "code2wav_enable_npu_graph": True,
+            "code2wav_max_npu_graphs": 32,
+        }
+
     def test_pipeline_exposes_full_and_async_payload_hooks(self) -> None:
         pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
         talker = pipeline.get_stage(1)

tests/platforms/npu/test_graph_tools.py

@@ -0,0 +1,333 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.platforms.npu.graph_tools import (
+    CapturedDeviceGraph,
+    NPUExactGraphRunner,
+)
+from vllm_omni.platforms.npu.models import minicpmo_4_5_code2wav as code2wav_patch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_captured_graph_replay_clones_persistent_outputs():
+    static_input = torch.zeros(1)
+    static_output = torch.zeros(1)
+
+    class _Graph:
+        def replay(self):
+            static_output.copy_(static_input * 2)
+
+    graph = CapturedDeviceGraph(
+        graph=_Graph(),
+        static_inputs=(static_input,),
+        static_outputs=(static_output,),
+    )
+
+    first = graph.replay((torch.tensor([2.0]),))[0]
+    second = graph.replay((torch.tensor([3.0]),))[0]
+
+    torch.testing.assert_close(first, torch.tensor([4.0]))
+    torch.testing.assert_close(second, torch.tensor([6.0]))
+    assert first.data_ptr() != second.data_ptr()
+
+
+def test_capture_uses_vllm_global_graph_pool(monkeypatch):
+    runner = NPUExactGraphRunner()
+    global_pool = object()
+    seen_pools = []
+    synchronizations = 0
+
+    class _Graph:
+        def replay(self):
+            pass
+
+    @contextmanager
+    def graph_context(graph, *, pool):
+        del graph
+        seen_pools.append(pool)
+        yield
+
+    def synchronize():
+        nonlocal synchronizations
+        synchronizations += 1
+
+    fake_npu = SimpleNamespace(
+        NPUGraph=_Graph,
+        graph=graph_context,
+        synchronize=synchronize,
+    )
+    monkeypatch.setattr(torch, "npu", fake_npu, raising=False)
+
+    import vllm.platforms
+
+    monkeypatch.setattr(
+        vllm.platforms,
+        "current_platform",
+        SimpleNamespace(get_global_graph_pool=lambda: global_pool),
+    )
+
+    captured = runner.capture(
+        (torch.tensor([2.0]),),
+        lambda value: (value + 1,),
+    )
+
+    assert seen_pools == [global_pool]
+    assert synchronizations == 2
+    torch.testing.assert_close(captured.static_outputs[0], torch.tensor([3.0]))
+
+
+def test_exact_signature_dispatch_captures_then_replays(monkeypatch):
+    runner = NPUExactGraphRunner(max_graphs=2)
+    replay_count = 0
+
+    class _FunctionalGraph:
+        def __init__(self, compute):
+            self.compute = compute
+
+        def replay(self, inputs):
+            nonlocal replay_count
+            replay_count += 1
+            return tuple(output.detach().clone() for output in self.compute(*inputs))
+
+    monkeypatch.setattr(runner, "_eligible", lambda inputs: True)
+    monkeypatch.setattr(
+        runner,
+        "capture",
+        lambda inputs, compute: _FunctionalGraph(compute),
+    )
+
+    def compute(value):
+        return (value * 2,)
+
+    first = runner.run("unit", (torch.tensor([2.0]),), (False,), compute)
+    second = runner.run("unit", (torch.tensor([3.0]),), (False,), compute)
+
+    torch.testing.assert_close(first[0], torch.tensor([4.0]))
+    torch.testing.assert_close(second[0], torch.tensor([6.0]))
+    assert runner.stats == {"captures": 1, "failed": 0, "hits": 1}
+    assert replay_count == 1
+
+
+def test_capture_failure_is_fatal_for_stage_process(monkeypatch):
+    runner = NPUExactGraphRunner(
+        component_name="test component",
+        disable_config_hint="disable the test graph",
+    )
+    captures = 0
+    computes = 0
+
+    monkeypatch.setattr(runner, "_eligible", lambda inputs: True)
+
+    def fail_capture(inputs, compute):
+        nonlocal captures
+        captures += 1
+        raise RuntimeError("unsupported op")
+
+    monkeypatch.setattr(runner, "capture", fail_capture)
+
+    def compute(value):
+        nonlocal computes
+        computes += 1
+        return (value + 1,)
+
+    with pytest.raises(RuntimeError, match="restart the stage process"):
+        runner.run("unit", (torch.tensor([1.0]),), (False,), compute)
+    with pytest.raises(RuntimeError, match="cannot continue"):
+        runner.run(
+            "different-shape",
+            (torch.tensor([1.0, 2.0]),),
+            (True,),
+            compute,
+        )
+
+    assert captures == 1
+    assert computes == 1
+    assert runner.stats == {"captures": 0, "failed": 1, "hits": 0}
+
+
+def test_code2wav_runtime_disables_internal_format_and_jit(monkeypatch):
+    monkeypatch.delenv("ASCEND_LAUNCH_BLOCKING", raising=False)
+    compile_modes = []
+    fake_npu = SimpleNamespace(
+        config=SimpleNamespace(allow_internal_format=True),
+        set_compile_mode=lambda **kwargs: compile_modes.append(kwargs),
+    )
+    monkeypatch.setattr(torch, "npu", fake_npu, raising=False)
+
+    code2wav_patch.prepare_code2wav_graph_runtime()
+
+    assert fake_npu.config.allow_internal_format is False
+    assert compile_modes == [{"jit_compile": False}]
+
+
+def test_code2wav_runtime_rejects_launch_blocking(monkeypatch):
+    monkeypatch.setenv("ASCEND_LAUNCH_BLOCKING", "1")
+
+    with pytest.raises(RuntimeError, match="ASCEND_LAUNCH_BLOCKING=1"):
+        code2wav_patch.prepare_code2wav_graph_runtime()
+
+
+@pytest.mark.parametrize(
+    ("enabled", "max_graphs", "expected_prepared"),
+    [
+        (True, 7, 1),
+        (False, 7, 0),
+        (True, 0, 0),
+    ],
+)
+def test_code2wav_patch_reads_stage_additional_config(
+    monkeypatch,
+    enabled,
+    max_graphs,
+    expected_prepared,
+):
+    prepared = 0
+
+    class _Backend:
+        speech_window = torch.zeros(1)
+        flow = SimpleNamespace(training=False)
+
+    model = SimpleNamespace(
+        backend=None,
+        vllm_config=SimpleNamespace(
+            additional_config={
+                "code2wav_enable_npu_graph": enabled,
+                "code2wav_max_npu_graphs": max_graphs,
+            }
+        ),
+    )
+
+    def build_backend(instance):
+        instance.backend = _Backend()
+
+    def prepare_runtime():
+        nonlocal prepared
+        prepared += 1
+
+    monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend)
+    monkeypatch.setattr(code2wav_patch, "prepare_code2wav_graph_runtime", prepare_runtime)
+    code2wav_patch._patched_build_backend(model)
+
+    assert prepared == expected_prepared
+    if expected_prepared:
+        graph_runner = code2wav_patch._backend_graph_runners[model.backend]
+        assert isinstance(graph_runner, NPUExactGraphRunner)
+        assert graph_runner.max_graphs == max_graphs
+    else:
+        assert model.backend not in code2wav_patch._backend_graph_runners
+
+
+@pytest.mark.parametrize("with_cache", [False, True])
+def test_code2wav_estimator_graph_dispatch_is_platform_owned(monkeypatch, with_cache):
+    calls = []
+
+    class _Backend:
+        _trt_stepper = None
+        _cfm_graph_wrapper = None
+
+    class _Runner:
+        def run(self, operation, inputs, constants, compute):
+            calls.append((operation, inputs, constants))
+            return compute(*inputs)
+
+    backend = _Backend()
+    runner = _Runner()
+    code2wav_patch._backend_graph_runners[backend] = runner
+
+    def graphable_estimator_step(
+        instance,
+        estimator,
+        *,
+        x,
+        mu,
+        time_embedding,
+        speakers,
+        cond,
+        cnn_cache,
+        att_cache,
+    ):
+        del instance, estimator, time_embedding, speakers, cond
+        cache_marker = x.new_tensor(0 if cnn_cache is None else 1)
+        return x + mu, cache_marker, cache_marker.clone()
+
+    monkeypatch.setattr(code2wav_patch, "_original_estimator_step", lambda *args, **kwargs: None)
+    monkeypatch.setattr(code2wav_patch, "_graphable_estimator_step", graphable_estimator_step)
+    value = torch.tensor([2.0])
+    cache = torch.tensor([1.0]) if with_cache else None
+    estimator = SimpleNamespace(t_embedder=lambda time: time + 1)
+    outputs = code2wav_patch._patched_estimator_step(
+        backend,
+        estimator,
+        x=value,
+        mu=value,
+        time=value,
+        speakers=value,
+        cond=value,
+        cnn_cache=cache,
+        att_cache=cache,
+    )
+
+    assert calls[0][0] == "cfm_estimator"
+    assert calls[0][2] == (with_cache,)
+    assert len(calls[0][1]) == (7 if with_cache else 5)
+    torch.testing.assert_close(calls[0][1][2], torch.tensor([[3.0]]))
+    torch.testing.assert_close(outputs[0], torch.tensor([4.0]))
+    torch.testing.assert_close(outputs[1], torch.tensor(1.0 if with_cache else 0.0))
+
+
+def test_code2wav_platform_wraps_flow_execution_context(monkeypatch):
+    entered = []
+
+    class _Backend:
+        pass
+
+    class _Context:
+        def __init__(self, device, require_math):
+            self.device = device
+            self.require_math = require_math
+
+        def __enter__(self):
+            entered.append((self.device, self.require_math))
+
+        def __exit__(self, exc_type, exc, traceback):
+            return False
+
+    backend = _Backend()
+    code2wav_patch._backend_graph_runners[backend] = NPUExactGraphRunner()
+    monkeypatch.setattr(
+        code2wav_patch,
+        "_flow_execution_context",
+        lambda device, *, require_math: _Context(device, require_math),
+    )
+    monkeypatch.setattr(
+        code2wav_patch,
+        "_original_setup_batch",
+        lambda instance, features, batch_size: ("setup", batch_size),
+    )
+    monkeypatch.setattr(
+        code2wav_patch,
+        "_original_decode_batch",
+        lambda instance, tokens, features, states, *, last_chunk, flush_encoder: (
+            "decode",
+            last_chunk,
+            flush_encoder,
+        ),
+    )
+    features = SimpleNamespace(speech_tokens=torch.zeros(1))
+
+    assert code2wav_patch._patched_setup_batch(backend, features, 3) == ("setup", 3)
+    assert code2wav_patch._patched_decode_batch(
+        backend,
+        torch.zeros(1),
+        features,
+        [],
+        last_chunk=True,
+    ) == ("decode", True, False)
+    assert entered == [(torch.device("cpu"), True), (torch.device("cpu"), True)]

tests/platforms/npu/test_minicpmo_code2wav_npugraph.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Ascend hardware validation for MiniCPM-o Code2Wav NPUGraph replay."""
+
+from math import prod
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
+    BatchedToken2Wav,
+)
+from vllm_omni.platforms.npu.graph_tools import NPUExactGraphRunner
+from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import (
+    _graphable_estimator_step,
+    prepare_code2wav_graph_runtime,
+)
+from vllm_omni.platforms.npu.models.step_audio2_token2wav import (
+    npu_token2wav_sdpa_context,
+)
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.omni,
+    *hardware_marks(res={"npu": "A3"}, num_cards=1),
+]
+
+
+class _Flow(nn.Module):
+    def __init__(self, estimator: nn.Module):
+        super().__init__()
+        self.decoder = SimpleNamespace(estimator=estimator)
+
+
+class _Token2Wav:
+    def __init__(self, estimator: nn.Module):
+        self.flow = _Flow(estimator).eval()
+        self.hift = nn.Identity().eval()
+        self.float16 = False
+        self.n_timesteps = 2
+        self.mel_cache_len = 1
+        self.source_cache_len = 2
+        self.speech_window = torch.ones(4, device="npu")
+
+
+def test_minicpmo_cfm_estimator_npugraph_matches_eager():
+    decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit")
+    prepare_code2wav_graph_runtime()
+
+    estimator = (
+        decoder_dit.DiT(
+            in_channels=6,
+            out_channels=2,
+            depth=2,
+            num_heads=2,
+            head_dim=8,
+            hidden_size=16,
+            mlp_ratio=2.0,
+        )
+        .npu()
+        .eval()
+    )
+    graph_runner = NPUExactGraphRunner(max_graphs=4)
+    adapter = BatchedToken2Wav(_Token2Wav(estimator))
+
+    def inputs(seed: int):
+        def make(shape: tuple[int, ...], offset: float):
+            values = torch.arange(
+                prod(shape),
+                device="npu",
+                dtype=torch.float32,
+            )
+            return torch.sin(values * 0.17 + seed + offset).reshape(shape)
+
+        return (
+            make((2, 2, 6), 0.0),
+            make((2, 2, 6), 0.5),
+            make((2, 1, 16), 1.0),
+            make((2, 1), 1.5),
+            make((2, 1, 6), 2.0),
+        )
+
+    def compute(*values, caches=None):
+        cnn_cache, att_cache = (None, None) if caches is None else caches
+        return _graphable_estimator_step(
+            adapter,
+            estimator,
+            x=values[0],
+            mu=values[1],
+            time_embedding=values[2],
+            speakers=values[3],
+            cond=values[4],
+            cnn_cache=cnn_cache,
+            att_cache=att_cache,
+        )
+
+    with torch.inference_mode(), npu_token2wav_sdpa_context(require_math=True):
+        warm_inputs = inputs(11)
+        warm_outputs = graph_runner.run(
+            "cfm_estimator",
+            warm_inputs,
+            (False,),
+            lambda *values: compute(*values),
+        )
+        replay_inputs = inputs(13)
+        expected = compute(*replay_inputs)
+        actual = graph_runner.run(
+            "cfm_estimator",
+            replay_inputs,
+            (False,),
+            lambda *values: compute(*values),
+        )
+        for eager, replayed in zip(expected, actual, strict=True):
+            torch.testing.assert_close(replayed, eager, rtol=1e-3, atol=1e-3)
+
+        cached_inputs = (*inputs(17), warm_outputs[1], warm_outputs[2])
+        graph_runner.run(
+            "cfm_estimator",
+            cached_inputs,
+            (True,),
+            lambda *values: compute(*values[:5], caches=(values[5], values[6])),
+        )
+        replay_cached_inputs = (
+            *inputs(19),
+            warm_outputs[1] + 0.1,
+            warm_outputs[2] + 0.1,
+        )
+        expected_cached = compute(
+            *replay_cached_inputs[:5],
+            caches=(replay_cached_inputs[5], replay_cached_inputs[6]),
+        )
+        actual_cached = graph_runner.run(
+            "cfm_estimator",
+            replay_cached_inputs,
+            (True,),
+            lambda *values: compute(*values[:5], caches=(values[5], values[6])),
+        )
+        for eager, replayed in zip(expected_cached, actual_cached, strict=True):
+            torch.testing.assert_close(replayed, eager, rtol=1e-3, atol=1e-3)
+        assert graph_runner.stats == {"captures": 2, "failed": 0, "hits": 2}

tests/platforms/npu/test_step_audio2_token2wav.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 import importlib.util
 import sys
 import types
+from contextlib import contextmanager
 from pathlib import Path
 from types import SimpleNamespace
 
@@ -194,3 +195,71 @@ def test_hift_patch_reports_incompatible_layout(monkeypatch: pytest.MonkeyPatch)
 
     with pytest.raises(TypeError, match=r"m_source\.l_sin_gen\._f02sine"):
         module.patch_step_audio2_hift_for_npu(SimpleNamespace())
+
+
+def test_sdpa_context_preserves_exceptions_from_token2wav_body(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    module = _load_patch_module(monkeypatch)
+    fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
+    fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None
+
+    @contextmanager
+    def fake_math_context(*, require_available=False):
+        del require_available
+        yield
+
+    fake_attention_patch.npu_math_sdpa_context = fake_math_context
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
+        fake_attention_patch,
+    )
+
+    with pytest.raises(RuntimeError, match="original graph failure"):
+        with module.npu_token2wav_sdpa_context():
+            raise RuntimeError("original graph failure")
+
+
+def test_sdpa_context_requires_math_for_npu_graph(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _load_patch_module(monkeypatch)
+    required = []
+    fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
+    fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None
+
+    @contextmanager
+    def fake_math_context(*, require_available=False):
+        required.append(require_available)
+        yield
+
+    fake_attention_patch.npu_math_sdpa_context = fake_math_context
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
+        fake_attention_patch,
+    )
+
+    with module.npu_token2wav_sdpa_context(require_math=True):
+        pass
+
+    assert required == [True]
+
+
+def test_sdpa_context_does_not_hide_patch_setup_failures(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _load_patch_module(monkeypatch)
+    fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
+
+    def fail_patch():
+        raise RuntimeError("attention patch failed")
+
+    fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = fail_patch
+    fake_attention_patch.npu_math_sdpa_context = contextmanager(lambda: iter((None,)))
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
+        fake_attention_patch,
+    )
+
+    with pytest.raises(RuntimeError, match="attention patch failed"):
+        with module.npu_token2wav_sdpa_context():
+            pass

vllm_omni/deploy/minicpmo_4_5.yaml

@@ -111,3 +111,9 @@ platforms:
         max_num_batched_tokens: 8192
         compilation_config:
           cudagraph_mode: PIECEWISE
+      - stage_id: 2
+        # Stage 2 stays outer-eager and captures only exact-shape CFM DiT
+        # estimator kernels through the NPU platform patch.
+        additional_config:
+          code2wav_enable_npu_graph: true
+          code2wav_max_npu_graphs: 32

vllm_omni/deploy/minicpmo_4_5_2gpu.yaml

@@ -93,3 +93,9 @@ platforms:
         max_num_batched_tokens: 8192
         compilation_config:
           cudagraph_mode: PIECEWISE
+      - stage_id: 2
+        # Stage 2 stays outer-eager and captures only exact-shape CFM DiT
+        # estimator kernels through the NPU platform patch.
+        additional_config:
+          code2wav_enable_npu_graph: true
+          code2wav_max_npu_graphs: 32

vllm_omni/deploy/minicpmo_4_5_3gpu.yaml

@@ -86,4 +86,20 @@ stages:
       max_tokens: 65536
       detokenize: true
 
-platforms: {}
+platforms:
+  npu:
+    stages:
+      - stage_id: 0
+        max_num_batched_tokens: 8192
+        compilation_config:
+          cudagraph_mode: PIECEWISE
+      - stage_id: 1
+        max_num_batched_tokens: 8192
+        compilation_config:
+          cudagraph_mode: PIECEWISE
+      - stage_id: 2
+        # Stage 2 stays outer-eager and captures only exact-shape CFM DiT
+        # estimator kernels through the NPU platform patch.
+        additional_config:
+          code2wav_enable_npu_graph: true
+          code2wav_max_npu_graphs: 32

vllm_omni/platforms/npu/__init__.py

@@ -1,6 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm_omni.platforms.npu.platform import NPUOmniPlatform
 
 __all__ = ["NPUOmniPlatform"]
+
+
+def __getattr__(name: str):
+    if name != "NPUOmniPlatform":
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+    globals()[name] = NPUOmniPlatform
+    return NPUOmniPlatform

vllm_omni/platforms/npu/graph_tools.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reusable exact-signature NPUGraph capture and replay helpers."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Protocol
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class _ReplayGraph(Protocol):
+    def replay(self) -> None: ...
+
+
+@dataclass
+class CapturedDeviceGraph:
+    graph: _ReplayGraph
+    static_inputs: tuple[torch.Tensor, ...]
+    static_outputs: tuple[torch.Tensor, ...]
+
+    def replay(self, inputs: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
+        with torch.inference_mode():
+            for static, current in zip(self.static_inputs, inputs, strict=True):
+                static.copy_(current)
+            self.graph.replay()
+            # Graph outputs are persistent and overwritten by the next replay.
+            # Clone before they become request-owned streaming cache entries.
+            return tuple(output.detach().clone() for output in self.static_outputs)
+
+
+def _tensor_signature(value: torch.Tensor) -> tuple[tuple[int, ...], str, str]:
+    return tuple(value.shape), str(value.dtype), str(value.device)
+
+
+class NPUExactGraphRunner:
+    """Capture and replay tensor-only functions for exact NPU signatures."""
+
+    def __init__(
+        self,
+        *,
+        max_graphs: int = 32,
+        component_name: str = "device graph",
+        disable_config_hint: str = "disable graph capture",
+    ) -> None:
+        self.max_graphs = max(0, int(max_graphs))
+        self.component_name = component_name
+        self.disable_config_hint = disable_config_hint
+        self._enabled = self.max_graphs > 0
+        self._graphs: dict[tuple[object, ...], CapturedDeviceGraph] = {}
+        self._failed_keys: set[tuple[object, ...]] = set()
+        self._hits = 0
+        self._graph_pool: object | None = None
+
+    @staticmethod
+    def is_supported() -> bool:
+        npu = getattr(torch, "npu", None)
+        return npu is not None and all(
+            hasattr(npu, name)
+            for name in (
+                "NPUGraph",
+                "graph",
+                "is_current_stream_capturing",
+                "synchronize",
+            )
+        )
+
+    @staticmethod
+    def _stream_is_capturing() -> bool:
+        npu = getattr(torch, "npu", None)
+        is_capturing = getattr(npu, "is_current_stream_capturing", None)
+        if not callable(is_capturing):
+            return False
+        try:
+            return bool(is_capturing())
+        except (RuntimeError, TypeError):
+            return False
+
+    def _eligible(self, inputs: tuple[torch.Tensor, ...]) -> bool:
+        return (
+            self._enabled
+            and bool(inputs)
+            and all(value.device.type == "npu" for value in inputs)
+            and self.is_supported()
+            and not self._stream_is_capturing()
+        )
+
+    @property
+    def stats(self) -> dict[str, int]:
+        return {
+            "captures": len(self._graphs),
+            "failed": len(self._failed_keys),
+            "hits": self._hits,
+        }
+
+    def capture(
+        self,
+        inputs: tuple[torch.Tensor, ...],
+        compute: Callable[..., tuple[torch.Tensor, ...]],
+    ) -> CapturedDeviceGraph:
+        npu = torch.npu
+        static_inputs = tuple(value.detach().clone() for value in inputs)
+        npu.synchronize()
+        graph = npu.NPUGraph()
+        if self._graph_pool is None:
+            from vllm.platforms import current_platform
+
+            self._graph_pool = current_platform.get_global_graph_pool()
+        with torch.inference_mode(), npu.graph(graph, pool=self._graph_pool):
+            static_outputs = compute(*static_inputs)
+        npu.synchronize()
+        return CapturedDeviceGraph(
+            graph=graph,
+            static_inputs=static_inputs,
+            static_outputs=static_outputs,
+        )
+
+    def run(
+        self,
+        operation: str,
+        inputs: tuple[torch.Tensor, ...],
+        constants: tuple[object, ...],
+        compute: Callable[..., tuple[torch.Tensor, ...]],
+    ) -> tuple[torch.Tensor, ...]:
+        if self._failed_keys:
+            raise RuntimeError(
+                f"{self.component_name} cannot continue after a failed NPUGraph capture; "
+                f"restart the stage process and {self.disable_config_hint} before retrying."
+            )
+        if not self._eligible(inputs):
+            return compute(*inputs)
+
+        key = (
+            operation,
+            constants,
+            tuple(_tensor_signature(value) for value in inputs),
+        )
+        graph = self._graphs.get(key)
+        if graph is not None:
+            self._hits += 1
+            if self._hits == 1:
+                logger.info("%s started NPUGraph replay", self.component_name)
+            return graph.replay(inputs)
+
+        # Prime lazy kernels and allocator state before capture. The next call
+        # with the same exact tensor signature replays this graph.
+        eager_outputs = compute(*inputs)
+        if len(self._graphs) >= self.max_graphs:
+            logger.warning_once(
+                "%s reached the %d-entry NPUGraph limit; new tensor shapes will use eager execution.",
+                self.component_name,
+                self.max_graphs,
+            )
+            return eager_outputs
+        try:
+            self._graphs[key] = self.capture(inputs, compute)
+        except Exception as exc:
+            self._failed_keys.add(key)
+            self._enabled = False
+            logger.exception(
+                "%s failed to capture NPUGraph for %s; the torch-npu "
+                "allocator/RNG capture state may be invalid and this stage "
+                "process must be restarted.",
+                self.component_name,
+                operation,
+            )
+            raise RuntimeError(
+                f"{self.component_name} NPUGraph capture failed for {operation}; "
+                f"restart the stage process. To run eagerly, {self.disable_config_hint}."
+            ) from exc
+        logger.info(
+            "%s captured NPUGraph %d/%d for %s",
+            self.component_name,
+            len(self._graphs),
+            self.max_graphs,
+            operation,
+        )
+        return eager_outputs

vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py

@@ -18,7 +18,7 @@ This module:
 from __future__ import annotations
 
 from collections.abc import Iterator
-from contextlib import contextmanager, nullcontext
+from contextlib import ExitStack, contextmanager
 
 import torch
 import torch.nn.functional as F
@@ -107,17 +107,20 @@ def _patched_attention_forward_chunk(
 
 
 @contextmanager
-def npu_math_sdpa_context() -> Iterator[None]:
+def npu_math_sdpa_context(*, require_available: bool = False) -> Iterator[None]:
     """Force SDPA MATH backend so Ascend does not call fused FA."""
-    try:
-        from torch.nn.attention import SDPBackend, sdpa_kernel
+    with ExitStack() as stack:
+        try:
+            from torch.nn.attention import SDPBackend, sdpa_kernel
 
-        with sdpa_kernel(SDPBackend.MATH):
-            yield
-    except Exception:
-        # Older torch / missing backend enum — just run as-is.
-        with nullcontext():
-            yield
+            stack.enter_context(sdpa_kernel(SDPBackend.MATH))
+        except Exception as exc:
+            if require_available:
+                raise RuntimeError(
+                    "MiniCPM-o Code2Wav NPU graph capture requires torch.nn.attention.sdpa_kernel(SDPBackend.MATH)"
+                ) from exc
+            logger.warning_once("MATH SDPA selection is unavailable; using the default Ascend SDPA backend.")
+        yield
 
 
 def _disable_upsample_encoder_compile() -> None:

vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inject MiniCPM-o Code2Wav NPUGraph acceleration on Ascend."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+from contextlib import nullcontext
+from weakref import WeakKeyDictionary
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.platforms.npu.graph_tools import NPUExactGraphRunner
+
+logger = init_logger(__name__)
+
+_PATCHED = False
+_original_build_backend = None
+_original_estimator_step = None
+_original_setup_batch = None
+_original_decode_batch = None
+_backend_graph_runners: WeakKeyDictionary[object, NPUExactGraphRunner] = WeakKeyDictionary()
+_ENABLE_KEY = "code2wav_enable_npu_graph"
+_MAX_GRAPHS_KEY = "code2wav_max_npu_graphs"
+
+
+def _config_bool(value: object, default: bool) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "on"}
+    return bool(value)
+
+
+def _graph_config(model: object) -> dict[str, object]:
+    config = getattr(getattr(model, "vllm_config", None), "additional_config", None)
+    return dict(config) if isinstance(config, Mapping) else {}
+
+
+def prepare_code2wav_graph_runtime() -> None:
+    """Select graph-capturable ACLNN kernels before Token2Wav is loaded."""
+    if os.environ.get("ASCEND_LAUNCH_BLOCKING") == "1":
+        raise RuntimeError(
+            "MiniCPM-o Code2Wav NPUGraph capture is incompatible with "
+            "ASCEND_LAUNCH_BLOCKING=1; unset it or set it to 0 before startup."
+        )
+    npu = torch.npu
+    npu.config.allow_internal_format = False
+    npu.set_compile_mode(jit_compile=False)
+    logger.info("Configured MiniCPM-o Code2Wav NPUGraph runtime (allow_internal_format=False, jit_compile=False)")
+
+
+def _flow_execution_context(device: torch.device, *, require_math: bool):
+    if device.type != "npu":
+        return nullcontext()
+    from vllm_omni.platforms.npu.models.step_audio2_token2wav import (
+        npu_token2wav_sdpa_context,
+    )
+
+    return npu_token2wav_sdpa_context(require_math=require_math)
+
+
+def _graphable_estimator_step(
+    backend,
+    estimator,
+    *,
+    x,
+    mu,
+    time_embedding,
+    speakers,
+    cond,
+    cnn_cache,
+    att_cache,
+):
+    """Run the CFM estimator body after host-backed timestep embedding."""
+    width = int(x.shape[-1])
+    speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width)
+    estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1)
+    cnn_out, att_out = backend._estimator_buffers(estimator, estimator_input, att_cache)
+    old_cnn = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks)
+    old_att = att_cache if att_cache is not None else [None] * len(estimator.blocks)
+    result = estimator.blocks_forward_chunk(
+        estimator_input,
+        time_embedding,
+        None,
+        old_cnn,
+        old_att,
+        cnn_out,
+        att_out,
+    )
+    return result, cnn_out, att_out
+
+
+def _patched_estimator_step(
+    self,
+    estimator,
+    *,
+    x,
+    mu,
+    time,
+    speakers,
+    cond,
+    cnn_cache,
+    att_cache,
+):
+    assert _original_estimator_step is not None
+    graph_runner = _backend_graph_runners.get(self)
+    if graph_runner is None or self._trt_stepper is not None or self._cfm_graph_wrapper is not None:
+        return _original_estimator_step(
+            self,
+            estimator,
+            x=x,
+            mu=mu,
+            time=time,
+            speakers=speakers,
+            cond=cond,
+            cnn_cache=cnn_cache,
+            att_cache=att_cache,
+        )
+    if (cnn_cache is None) != (att_cache is None):
+        raise ValueError("estimator CNN and attention caches must both be present or absent")
+
+    # The upstream embedder creates a frequency tensor on the host. Keep it
+    # outside capture while retaining the tensor-only estimator body in graph.
+    time_embedding = estimator.t_embedder(time).unsqueeze(1)
+    if cnn_cache is None:
+        return graph_runner.run(
+            "cfm_estimator",
+            (x, mu, time_embedding, speakers, cond),
+            (False,),
+            lambda step_x, step_mu, step_time, step_speakers, step_cond: _graphable_estimator_step(
+                self,
+                estimator,
+                x=step_x,
+                mu=step_mu,
+                time_embedding=step_time,
+                speakers=step_speakers,
+                cond=step_cond,
+                cnn_cache=None,
+                att_cache=None,
+            ),
+        )
+
+    return graph_runner.run(
+        "cfm_estimator",
+        (x, mu, time_embedding, speakers, cond, cnn_cache, att_cache),
+        (True,),
+        lambda step_x, step_mu, step_time, step_speakers, step_cond, step_cnn, step_att: _graphable_estimator_step(
+            self,
+            estimator,
+            x=step_x,
+            mu=step_mu,
+            time_embedding=step_time,
+            speakers=step_speakers,
+            cond=step_cond,
+            cnn_cache=step_cnn,
+            att_cache=step_att,
+        ),
+    )
+
+
+def _patched_setup_batch(self, features, batch_size):
+    assert _original_setup_batch is not None
+    with _flow_execution_context(
+        features.speech_tokens.device,
+        require_math=self in _backend_graph_runners,
+    ):
+        return _original_setup_batch(self, features, batch_size)
+
+
+def _patched_decode_batch(
+    self,
+    tokens,
+    features,
+    states,
+    *,
+    last_chunk,
+    flush_encoder=False,
+):
+    assert _original_decode_batch is not None
+    with _flow_execution_context(
+        tokens.device,
+        require_math=self in _backend_graph_runners,
+    ):
+        return _original_decode_batch(
+            self,
+            tokens,
+            features,
+            states,
+            last_chunk=last_chunk,
+            flush_encoder=flush_encoder,
+        )
+
+
+def _patched_build_backend(self) -> None:
+    if self.backend is not None:
+        return
+
+    config = _graph_config(self)
+    max_graphs = max(0, int(config.get(_MAX_GRAPHS_KEY, 32)))
+    graph_enabled = max_graphs > 0 and _config_bool(config.get(_ENABLE_KEY), False)
+    if graph_enabled:
+        # NPUOmniPlatform enables internal format for quantized LLM kernels.
+        # Code2Wav uses regular convolution kernels that must remain in the
+        # graph-capturable ACLNN path.
+        prepare_code2wav_graph_runtime()
+
+    assert _original_build_backend is not None
+    _original_build_backend(self)
+
+    graph_runner = None
+    if graph_enabled:
+        graph_runner = NPUExactGraphRunner(
+            max_graphs=max_graphs,
+            component_name="MiniCPM-o Code2Wav",
+            disable_config_hint=(
+                "set platforms.npu.stages[stage_id=2].additional_config.code2wav_enable_npu_graph=false"
+            ),
+        )
+        if self.backend.speech_window.device.type == "npu" and not graph_runner.is_supported():
+            raise RuntimeError(
+                "MiniCPM-o Code2Wav NPUGraph capture requires torch.npu "
+                "NPUGraph, graph, is_current_stream_capturing, and synchronize APIs."
+            )
+        if self.backend.flow.training:
+            raise ValueError("MiniCPM-o Code2Wav NPUGraph capture requires flow.eval()")
+        _backend_graph_runners[self.backend] = graph_runner
+
+    if graph_enabled:
+        logger.info(
+            "MiniCPM-o Code2Wav NPUGraph replay enabled (max_graphs=%d)",
+            max_graphs,
+        )
+
+
+def apply_minicpmo_4_5_code2wav_patch() -> None:
+    """Patch the generic Code2Wav backend builder with Ascend acceleration."""
+    global _PATCHED, _original_build_backend
+    global _original_decode_batch, _original_estimator_step, _original_setup_batch
+    if _PATCHED:
+        return
+
+    from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
+        BatchedToken2Wav,
+    )
+    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import (
+        MiniCPMO45Code2Wav,
+    )
+
+    _original_build_backend = MiniCPMO45Code2Wav._build_backend
+    _original_estimator_step = BatchedToken2Wav._estimator_step
+    _original_setup_batch = BatchedToken2Wav.setup_batch
+    _original_decode_batch = BatchedToken2Wav.decode_batch
+
+    MiniCPMO45Code2Wav._build_backend = _patched_build_backend  # type: ignore[method-assign]
+    BatchedToken2Wav._estimator_step = _patched_estimator_step  # type: ignore[method-assign]
+    BatchedToken2Wav.setup_batch = _patched_setup_batch  # type: ignore[method-assign]
+    BatchedToken2Wav.decode_batch = _patched_decode_batch  # type: ignore[method-assign]
+    _PATCHED = True
+    logger.debug("Applied NPU patch for MiniCPM-o 4.5 Code2Wav")

vllm_omni/platforms/npu/models/step_audio2_token2wav.py

@@ -13,7 +13,7 @@ Ascend-specific workarounds that must not live in the shared GPU model file:
 from __future__ import annotations
 
 from collections.abc import Iterator
-from contextlib import contextmanager, nullcontext
+from contextlib import contextmanager
 from types import MethodType
 
 import numpy as np
@@ -113,20 +113,16 @@ def patch_step_audio2_hift_for_npu(hift: torch.nn.Module) -> None:
 
 
 @contextmanager
-def npu_token2wav_sdpa_context() -> Iterator[None]:
+def npu_token2wav_sdpa_context(*, require_math: bool = False) -> Iterator[None]:
     """Expand CosyVoice masks + force MATH SDPA to avoid FA 161001."""
-    try:
-        from vllm_omni.platforms.npu.models.cosyvoice2_dit_attn import (
-            apply_cosyvoice2_dit_attn_npu_patch,
-            npu_math_sdpa_context,
-        )
+    from vllm_omni.platforms.npu.models.cosyvoice2_dit_attn import (
+        apply_cosyvoice2_dit_attn_npu_patch,
+        npu_math_sdpa_context,
+    )
 
-        apply_cosyvoice2_dit_attn_npu_patch()
-        with npu_math_sdpa_context():
-            yield
-    except Exception:
-        with nullcontext():
-            yield
+    apply_cosyvoice2_dit_attn_npu_patch()
+    with npu_math_sdpa_context(require_available=require_math):
+        yield
 
 
 def _patched_ensure_models_loaded(self) -> None:

vllm_omni/platforms/npu/platform.py

@@ -36,6 +36,9 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         from vllm_ascend.utils import adapt_patch
 
         from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches
+        from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import (
+            apply_minicpmo_4_5_code2wav_patch,
+        )
         from vllm_omni.platforms.npu.models.minimax_h3 import apply_minimax_h3_qwen3vl_patch
         from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import (
             apply_qwen3_tts_code2wav_patch,
@@ -46,6 +49,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
         adapt_patch(is_global_patch=True)
         apply_minimax_h3_qwen3vl_patch()
+        apply_minicpmo_4_5_code2wav_patch()
         apply_qwen3_tts_code2wav_patch()
         apply_qwen3_tts_tokenizer_v2_patch()
         apply_310p_patches()