From ae3880f0e39df4ba7376481ae78d9269fc42e09c Mon Sep 17 00:00:00 2001 From: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com> Date: Mon, 21 Sep 2026 06:38:23 +0800 Subject: [PATCH] [Performance] Capture MOSS-TTS codec streaming decode with NPUGraph (#7280) Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com> Co-authored-by: amy-why-3459 Co-authored-by: Gao Han --- .../npu/test_moss_tts_codec_npugraph.py | 205 ++++++++++++++++ vllm_omni/deploy/moss_tts_local.yaml | 11 + .../moss_tts/modeling_moss_tts_codec.py | 34 ++- .../moss_tts_streaming_decode_wrapper.py | 220 ++++++++++++++++++ 4 files changed, 460 insertions(+), 10 deletions(-) create mode 100644 tests/platforms/npu/test_moss_tts_codec_npugraph.py create mode 100644 vllm_omni/platforms/npu/models/moss_tts_streaming_decode_wrapper.py diff --git a/tests/platforms/npu/test_moss_tts_codec_npugraph.py b/tests/platforms/npu/test_moss_tts_codec_npugraph.py new file mode 100644 index 000000000..dee0f5e4d --- /dev/null +++ b/tests/platforms/npu/test_moss_tts_codec_npugraph.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Ascend hardware validation for Moss-TTS codec streaming NPUGraph replay. + +Mirrors ``test_codec_cudagraph.py`` (non-streaming CUDA wrapper) but exercises +the streaming :class:`NPUGraphStreamingDecoderWrapper` which captures +``decode_streaming_tensors`` per ``(B_bucket, exact_T)`` signature. + +The synthetic codec is stateless — graph vs eager parity is verified for the +same inputs. State management (RingKVCache in-place updates) is covered by +the E2E integration test (correct audio output confirmed on 910B). +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +pytest.importorskip("vllm_ascend") + +from tests.helpers.mark import hardware_marks +from vllm_omni.platforms.npu.models.moss_tts_streaming_decode_wrapper import ( + NPUGraphStreamingDecoderWrapper, +) + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.omni, + *hardware_marks(res={"npu": "A3"}, num_cards=1), +] + +DEVICE = torch.device("npu:0") +NQ = 4 +UPSAMPLE = 4 +HIDDEN = 32 + + +class SyntheticStreamingCodec(nn.Module): + """Minimal stateless stand-in for MossAudioTokenizerModel streaming decode. + + Exposes ``decode_streaming_tensors`` + ``reset_decoder_state_slots`` with + the same interface the NPU wrapper expects. + """ + + def __init__(self): + super().__init__() + self.downsample_rate = UPSAMPLE + self.embed = nn.Conv1d(NQ, HIDDEN, kernel_size=3, padding=1) + self.conv = nn.Conv1d(HIDDEN, HIDDEN, kernel_size=3, padding=1) + self.up = nn.ConvTranspose1d(HIDDEN, 1, kernel_size=UPSAMPLE, stride=UPSAMPLE) + + def decode_streaming_tensors(self, codes, lengths, slot_ids, valid_rows): + """codes [NQ,B,T] lengths [B] slot_ids [B] valid [B] -> audio [B,1,T*up].""" + nq, b, t = codes.shape + x = codes.permute(1, 0, 2).to(dtype=self.embed.weight.dtype) + x = torch.relu(self.embed(x)) + x = torch.relu(self.conv(x)) + audio = self.up(x) + audio_lengths = lengths * self.downsample_rate + return audio, audio_lengths + + def reset_decoder_state_slots(self, slot_ids): + pass + + +@pytest.fixture(scope="module") +def codec(): + torch.manual_seed(0) + return SyntheticStreamingCodec().to(DEVICE).eval() + + +@pytest.fixture(scope="module") +def wrapper(codec): + w = NPUGraphStreamingDecoderWrapper( + codec=codec, + state_capacity=16, + batch_sizes=[1, 2, 4], + frame_sizes=[4, 8], + num_quantizers=NQ, + vllm_config=SimpleNamespace(), + ) + w.warmup(DEVICE) + return w + + +def _codes(b, t, device=DEVICE): + return torch.randint(0, 64, (NQ, b, t), dtype=torch.long, device=device) + + +def _eager(codec, codes, slot_ids): + b = codes.shape[1] + t = codes.shape[2] + lengths = torch.full((b,), t, dtype=torch.long, device=codes.device) + valid = torch.ones(b, dtype=torch.bool, device=codes.device) + with torch.no_grad(): + return codec.decode_streaming_tensors(codes, lengths, slot_ids, valid) + + +# --------------------------------------------------------------------------- +# 1. Exact-size decode — graph replay must match eager +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("b,t", [(1, 4), (2, 8), (4, 4)]) +def test_exact_size_replay_matches_eager(codec, wrapper, b, t): + slots = torch.arange(b, dtype=torch.long, device=DEVICE) + codes = _codes(b, t) + ref = _eager(codec, codes, slots) + with torch.no_grad(): + out = wrapper.decode(codes, slots) + assert out is not None + audio, _, actual_b = out + assert actual_b == b + torch.testing.assert_close(audio[:b], ref[0][:b], atol=1e-3, rtol=1e-3) + + +# --------------------------------------------------------------------------- +# 2. Padded batch — B < bucket or T < frame_size +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("b,t", [(1, 5), (2, 4), (3, 6)]) +def test_padded_batch_interior_matches(codec, wrapper, b, t): + """Padded batches must match eager run on the same padded input. + + The graph wrapper zero-pads codes to the next frame_size bucket. We run + eager on the same padded codes so both paths see identical input. + """ + slots = torch.arange(b, dtype=torch.long, device=DEVICE) + codes = _codes(b, t) + graph_frame_size = next(fs for fs in wrapper.frame_sizes if fs >= t) + padded_codes = torch.zeros(NQ, b, graph_frame_size, dtype=torch.long, device=DEVICE) + padded_codes[:, :, :t] = codes + lengths = torch.full((b,), t, dtype=torch.long, device=DEVICE) + valid = torch.ones(b, dtype=torch.bool, device=DEVICE) + with torch.no_grad(): + ref = codec.decode_streaming_tensors(padded_codes, lengths, slots, valid) + with torch.no_grad(): + out = wrapper.decode(codes, slots, allow_frame_padding=True) + assert out is not None + audio, audio_lengths, actual_b = out + assert actual_b == b + expected_len = t * UPSAMPLE + for row in range(b): + assert int(audio_lengths[row]) == expected_len + torch.testing.assert_close(audio[row], ref[0][row], atol=1e-3, rtol=1e-3) + + +# --------------------------------------------------------------------------- +# 3. Terminal tail — finished request with T < frame_size, padded +# --------------------------------------------------------------------------- + + +def test_terminal_tail_matches_eager(codec, wrapper): + """Terminal tail (T=3 padded to frame_size=4) must match eager on the + same padded input.""" + slots = torch.tensor([0], dtype=torch.long, device=DEVICE) + t = 3 + codes = _codes(1, t) + graph_frame_size = next(fs for fs in wrapper.frame_sizes if fs >= t) + padded_codes = torch.zeros(NQ, 1, graph_frame_size, dtype=torch.long, device=DEVICE) + padded_codes[:, :, :t] = codes + lengths = torch.tensor([t], dtype=torch.long, device=DEVICE) + valid = torch.ones(1, dtype=torch.bool, device=DEVICE) + with torch.no_grad(): + ref = codec.decode_streaming_tensors(padded_codes, lengths, slots, valid) + with torch.no_grad(): + out = wrapper.decode(codes, slots, allow_frame_padding=True) + assert out is not None + audio, audio_lengths, _ = out + assert int(audio_lengths[0]) == t * UPSAMPLE + torch.testing.assert_close(audio[0], ref[0][0], atol=1e-3, rtol=1e-3) + + +# --------------------------------------------------------------------------- +# 4. Output not aliased across replays (static buffer clone) +# --------------------------------------------------------------------------- + + +def test_output_not_aliased(codec, wrapper): + """wrapper.decode returns cloned audio; a later replay must not overwrite + a previously returned tensor.""" + slot = torch.tensor([0], dtype=torch.long, device=DEVICE) + with torch.no_grad(): + out1 = wrapper.decode(_codes(1, 4), slot) + saved = out1[0].clone() + out2 = wrapper.decode(_codes(1, 4), slot) + assert out1 is not None and out2 is not None + torch.testing.assert_close(out1[0], saved, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 5. Unsupported device returns None (graceful skip) +# --------------------------------------------------------------------------- + + +def test_unsupported_device_returns_none(wrapper): + """decode on a non-NPU tensor must return None (no graph path).""" + if not torch.cuda.is_available(): + pytest.skip("CUDA not available for cross-device test") + codes = _codes(1, 4, device=torch.device("cuda:0")) + slots = torch.tensor([0], dtype=torch.long, device=codes.device) + out = wrapper.decode(codes, slots) + assert out is None diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index 5f0bf9a14..ad16a6dea 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -64,3 +64,14 @@ stages: max_tokens: 65536 seed: 42 repetition_penalty: 1.0 + +# Stage 1 NPU override: PIECEWISE cudagraph for NPUGraph capture; set enforce_eager: true to disable. +platforms: + npu: + stages: + - stage_id: 1 + compilation_config: + cudagraph_mode: PIECEWISE + # Re-declared: _apply_platform_overrides() wholesale-replaces compilation_config. + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64] + cudagraph_num_of_warmups: 1 diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index fbefe3a6a..6fd097d71 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -63,7 +63,7 @@ class _MossCodecStreamSession: self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None batch_sizes = sorted({int(size) for size in (graph_batch_sizes or []) if 0 < int(size) <= self._state_capacity}) frame_sizes = sorted({int(size) for size in (graph_frame_sizes or []) if int(size) > 0}) - scratch_capacity = max(batch_sizes, default=0) if self._device.type == "cuda" else 0 + scratch_capacity = max(batch_sizes, default=0) if self._device.type in ("cuda", "npu") else 0 self._total_state_capacity = self._state_capacity + scratch_capacity self._state_slot_ids = torch.arange( self._total_state_capacity, @@ -76,15 +76,29 @@ class _MossCodecStreamSession: raise RuntimeError("The streaming codec does not implement a decoder state pool.") with torch.no_grad(): initialize_state_pool(self._state_capacity, scratch_capacity) - if batch_sizes and frame_sizes and self._device.type == "cuda": - self._cudagraph_wrapper = CUDAGraphStreamingDecoderWrapper( - codec, - state_capacity=self._state_capacity, - batch_sizes=batch_sizes, - frame_sizes=frame_sizes, - num_quantizers=self._n_vq, - vllm_config=vllm_config, - ) + if batch_sizes and frame_sizes and self._device.type in ("cuda", "npu"): + if self._device.type == "npu": + from vllm_omni.platforms.npu.models.moss_tts_streaming_decode_wrapper import ( + NPUGraphStreamingDecoderWrapper, + ) + + self._cudagraph_wrapper = NPUGraphStreamingDecoderWrapper( + codec, + state_capacity=self._state_capacity, + batch_sizes=batch_sizes, + frame_sizes=frame_sizes, + num_quantizers=self._n_vq, + vllm_config=vllm_config, + ) + else: + self._cudagraph_wrapper = CUDAGraphStreamingDecoderWrapper( + codec, + state_capacity=self._state_capacity, + batch_sizes=batch_sizes, + frame_sizes=frame_sizes, + num_quantizers=self._n_vq, + vllm_config=vllm_config, + ) self._cudagraph_wrapper.warmup(self._device) self.reset_slots(list(range(self._state_capacity + scratch_capacity))) if not self._cudagraph_wrapper.is_ready: diff --git a/vllm_omni/platforms/npu/models/moss_tts_streaming_decode_wrapper.py b/vllm_omni/platforms/npu/models/moss_tts_streaming_decode_wrapper.py new file mode 100644 index 000000000..5928de4bb --- /dev/null +++ b/vllm_omni/platforms/npu/models/moss_tts_streaming_decode_wrapper.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MOSS-TTS codec streaming NPUGraph acceleration on Ascend. + +Mirrors :class:`CUDAGraphStreamingDecoderWrapper` but uses +:class:`NPUExactGraphRunner` for capture/replay. Graph capture is gated +by ``enforce_eager: false`` → ``cudagraph_capture_sizes`` → batch sizes; +set ``enforce_eager: true`` on Stage 1 to disable. +""" + +from __future__ import annotations + +import os +import time + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from vllm_omni.platforms.npu.graph_tools import NPUExactGraphRunner + +logger = init_logger(__name__) + + +def _prepare_moss_codec_graph_runtime() -> None: + """Select graph-capturable ACLNN kernels before MOSS-TTS codec capture. + + Mirrors ``prepare_code2wav_graph_runtime`` from minicpmo_4_5_code2wav: + ``allow_internal_format=False`` + ``jit_compile=False`` are required for + NPUGraph capture of the codec's conv/attention kernels, but must only be + set when graph capture is actually enabled. + """ + if os.environ.get("ASCEND_LAUNCH_BLOCKING") == "1": + raise RuntimeError( + "MOSS-TTS codec 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 MOSS-TTS codec NPUGraph runtime (allow_internal_format=False, jit_compile=False)") + + +class NPUGraphStreamingDecoderWrapper: + """Replay streaming decode graphs keyed by ``(B_bucket, exact_T)`` on NPU. + + Delegates capture/replay to :class:`NPUExactGraphRunner`. Bucket selection + (batch + frame padding) mirrors the CUDA wrapper so + ``_MossCodecStreamSession.step`` can use either interchangeably. Codec + state is updated in-place during replay via ``state_slot_ids``; scratch + slots isolate padding rows. + + Fatal NPUGraph failures propagate as ``RuntimeError`` (eager fallback is + unsafe after corrupt capture state); restart with ``enforce_eager: true`` + on Stage 1 to recover. + """ + + def __init__( + self, + codec: nn.Module, + *, + state_capacity: int, + batch_sizes: list[int], + frame_sizes: list[int], + num_quantizers: int, + vllm_config: VllmConfig, + ) -> None: + self.codec = codec + self.state_capacity = int(state_capacity) + self.batch_sizes = sorted({int(size) for size in batch_sizes if 0 < int(size) <= state_capacity}) + self.frame_sizes = sorted({int(size) for size in frame_sizes if int(size) > 0}) + self.num_quantizers = int(num_quantizers) + + # vllm_config is accepted for interface parity but unused; graph + # capture is gated by enforce_eager on Stage 1. + del vllm_config + _prepare_moss_codec_graph_runtime() + max_graphs = len(self.batch_sizes) * len(self.frame_sizes) + self._graph_runner: NPUExactGraphRunner | None = NPUExactGraphRunner( + max_graphs=max_graphs, + component_name="MOSS-TTS Codec", + disable_config_hint="set enforce_eager: true on Stage 1", + ) + if not self._graph_runner.is_supported(): + logger.warning("MOSS-TTS codec NPUGraph is not supported on this platform; falling back to eager.") + self._graph_runner = None + + @property + def is_ready(self) -> bool: + return self._graph_runner is not None + + @property + def scratch_capacity(self) -> int: + return max(self.batch_sizes, default=0) + + @torch.no_grad() + def warmup(self, device: torch.device) -> None: + if self._graph_runner is None or not self.batch_sizes or not self.frame_sizes: + return + if device.type != "npu": + return + + capture_keys = sorted( + ((b, t) for b in self.batch_sizes for t in self.frame_sizes), + reverse=True, + ) + logger.info( + "MOSS-TTS streaming decoder NPUGraph warmup: n_vq=%d (B,T)=%s", + self.num_quantizers, + list(reversed(capture_keys)), + ) + start_s = time.perf_counter() + for batch_size, frame_size in capture_keys: + self._warmup_bucket(batch_size, frame_size, device) + logger.info( + " Captured MOSS-TTS streaming decoder NPUGraph for (B,T)=(%d,%d)", + batch_size, + frame_size, + ) + + stats = self._graph_runner.stats if self._graph_runner else {} + captured = stats.get("captures", 0) + logger.info( + "MOSS-TTS streaming decoder NPUGraph warmup complete: %d/%d captured in %.1f ms (stats=%s)", + captured, + len(capture_keys), + (time.perf_counter() - start_s) * 1000.0, + stats, + ) + + @torch.no_grad() + def _warmup_bucket(self, batch_size: int, frame_size: int, device: torch.device) -> None: + """Trigger first-run capture for one (B,T) bucket using scratch slots. + + Feeds scratch slot ids + ``valid_rows=False`` so capture-side state + writes land in scratch slots, then resets them. + """ + assert self._graph_runner is not None + codes = torch.zeros( + self.num_quantizers, + batch_size, + frame_size, + dtype=torch.long, + device=device, + ) + lengths = torch.zeros(batch_size, dtype=torch.long, device=device) + scratch_slots = self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=device) + valid_rows = torch.zeros(batch_size, dtype=torch.bool, device=device) + + self._graph_runner.run( + "moss_codec_decode", + (codes, lengths, scratch_slots, valid_rows), + (batch_size, frame_size), + lambda c, lengths, s, v: self.codec.decode_streaming_tensors(c, lengths, s, v), + ) + self.codec.reset_decoder_state_slots(scratch_slots) + + def _select_batch_size(self, actual_batch_size: int) -> int | None: + return next((size for size in self.batch_sizes if size >= actual_batch_size), None) + + def _select_frame_size(self, actual_frame_size: int, allow_padding: bool) -> int | None: + if actual_frame_size in self.frame_sizes: + return actual_frame_size + if not allow_padding: + return None + return next((size for size in self.frame_sizes if size >= actual_frame_size), None) + + @torch.no_grad() + def decode( + self, + codes: torch.Tensor, + state_slot_ids: torch.Tensor, + *, + allow_frame_padding: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, int] | None: + if self._graph_runner is None: + return None + if codes.device.type != "npu": + return None + + n_vq, actual_batch_size, frame_size = codes.shape + if int(n_vq) != self.num_quantizers or state_slot_ids.shape != (actual_batch_size,): + return None + + batch_size = self._select_batch_size(int(actual_batch_size)) + if batch_size is None: + return None + graph_frame_size = self._select_frame_size(int(frame_size), allow_frame_padding) + if graph_frame_size is None: + return None + + device = codes.device + # Padding rows use scratch slots + valid_rows=False to isolate state. + padded_codes = torch.zeros( + self.num_quantizers, + batch_size, + graph_frame_size, + dtype=torch.long, + device=device, + ) + padded_codes[:, :actual_batch_size, :frame_size].copy_(codes) + lengths = torch.zeros(batch_size, dtype=torch.long, device=device) + lengths[:actual_batch_size].fill_(int(frame_size)) + padded_slots = self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=device) + padded_slots[:actual_batch_size].copy_(state_slot_ids) + valid_rows = torch.zeros(batch_size, dtype=torch.bool, device=device) + valid_rows[:actual_batch_size].fill_(True) + + audio, audio_lengths = self._graph_runner.run( + "moss_codec_decode", + (padded_codes, lengths, padded_slots, valid_rows), + (batch_size, graph_frame_size), + lambda c, lengths, s, v: self.codec.decode_streaming_tensors(c, lengths, s, v), + ) + + return audio, audio_lengths, int(actual_batch_size) + + +__all__ = ["NPUGraphStreamingDecoderWrapper"]