[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 <wuhaiyan17@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

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

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

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:

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