From 5d09cf27a98bb104506ee842ca81e0e76e47dc92 Mon Sep 17 00:00:00 2001 From: Zhou248 <62542847+Zhou248@users.noreply.github.com> Date: Mon, 17 Aug 2026 15:09:51 +0800 Subject: [PATCH] [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 | 11 + recipes/OpenBMB/MiniCPM-o-4_5.md | 39 +- .../models/minicpmo_4_5/test_pipeline.py | 25 ++ tests/platforms/npu/test_graph_tools.py | 333 ++++++++++++++++++ .../npu/test_minicpmo_code2wav_npugraph.py | 143 ++++++++ .../npu/test_step_audio2_token2wav.py | 69 ++++ vllm_omni/deploy/minicpmo_4_5.yaml | 6 + vllm_omni/deploy/minicpmo_4_5_2gpu.yaml | 6 + vllm_omni/deploy/minicpmo_4_5_3gpu.yaml | 18 +- vllm_omni/platforms/npu/__init__.py | 15 +- vllm_omni/platforms/npu/graph_tools.py | 183 ++++++++++ .../npu/models/cosyvoice2_dit_attn.py | 23 +- .../npu/models/minicpmo_4_5_code2wav.py | 262 ++++++++++++++ .../npu/models/step_audio2_token2wav.py | 22 +- vllm_omni/platforms/npu/platform.py | 4 + 15 files changed, 1129 insertions(+), 30 deletions(-) create mode 100644 tests/platforms/npu/test_graph_tools.py create mode 100644 tests/platforms/npu/test_minicpmo_code2wav_npugraph.py create mode 100644 vllm_omni/platforms/npu/graph_tools.py create mode 100644 vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py diff --git a/.buildkite/npu/test-npu-ready.yml b/.buildkite/npu/test-npu-ready.yml index f23ba7120..dbb827137 100644 --- a/.buildkite/npu/test-npu-ready.yml +++ b/.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 diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md index 99b064310..86f0883f7 100644 --- a/recipes/OpenBMB/MiniCPM-o-4_5.md +++ b/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) diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 148f28b16..bbc011284 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/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) diff --git a/tests/platforms/npu/test_graph_tools.py b/tests/platforms/npu/test_graph_tools.py new file mode 100644 index 000000000..7ec36dfe9 --- /dev/null +++ b/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)] diff --git a/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py b/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py new file mode 100644 index 000000000..1c35c37aa --- /dev/null +++ b/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} diff --git a/tests/platforms/npu/test_step_audio2_token2wav.py b/tests/platforms/npu/test_step_audio2_token2wav.py index ae86f92d9..23baa2c3e 100644 --- a/tests/platforms/npu/test_step_audio2_token2wav.py +++ b/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 diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index bc34bfe9f..04ddad608 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/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 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index 2a4abed42..5b43206b8 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/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 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index b34673910..edbde0c5d 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/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 diff --git a/vllm_omni/platforms/npu/__init__.py b/vllm_omni/platforms/npu/__init__.py index fb95f7b29..74dfc419c 100644 --- a/vllm_omni/platforms/npu/__init__.py +++ b/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 diff --git a/vllm_omni/platforms/npu/graph_tools.py b/vllm_omni/platforms/npu/graph_tools.py new file mode 100644 index 000000000..d943d1584 --- /dev/null +++ b/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 diff --git a/vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py b/vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py index 7e1d5d540..7cf280a00 100644 --- a/vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py +++ b/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: diff --git a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py new file mode 100644 index 000000000..c6effe6f8 --- /dev/null +++ b/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") diff --git a/vllm_omni/platforms/npu/models/step_audio2_token2wav.py b/vllm_omni/platforms/npu/models/step_audio2_token2wav.py index 1641fee27..1999a08a0 100644 --- a/vllm_omni/platforms/npu/models/step_audio2_token2wav.py +++ b/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: diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index b17f447c4..38c35f1f7 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/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()