[Hardware][Ascend][Model] Support MOSS-TTS-Nano on Ascend A2 (#7192)

Signed-off-by: Big2Wheel <1404207590@qq.com>
2026-09-14 15:25:17 -07:00
1 parent c68eadbbe0
commit 8e088c7725
@@ -87,7 +87,7 @@ th {
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | ✅︎ | — |
 | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — |
-| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — |
+| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | ✅︎ | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/OpenMOSS/MOSS-TTS-Nano-NPU.md) |
 | `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — |
 | `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | | — |
 | `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | — |
@@ -0,0 +1,128 @@
+# MOSS-TTS-Nano on Ascend NPU
+
+## Summary
+
+- Model: `OpenMOSS-Team/MOSS-TTS-Nano`
+- Task: multilingual text-to-speech with a reference voice
+- Hardware: one Ascend 910B3 64GB NPU
+- Output: 48 kHz mono audio
+- Serving API: OpenAI-compatible `/v1/audio/speech`
+
+MOSS-TTS-Nano combines a small autoregressive language model with
+`OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano`. Both components run on the same
+NPU. The codec is a separate checkpoint and must be available even when the
+main model has already been downloaded.
+
+## Environment
+
+Use an aligned vLLM, vLLM-Ascend, and vLLM-Omni environment. The following
+configuration was validated:
+
+- Image: `quay.nju.edu.cn/ascend/vllm-omni:v0.28.0`
+- Ascend device: 910B3 64GB
+- CANN: 9.1.0
+- vLLM: 0.28.0
+- Execution mode: eager
+
+See the [NPU installation guide](../../docs/getting_started/installation/npu.md)
+for the complete container device and driver mounts.
+
+## Download checkpoints
+
+```bash
+export HF_HOME=/models/huggingface-cache
+
+hf download OpenMOSS-Team/MOSS-TTS-Nano
+
+# This dependency is loaded from its model ID by the MOSS-TTS-Nano config.
+hf download OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano
+```
+
+For an air-gapped server, populate `HF_HOME` with the audio-tokenizer snapshot
+before startup and set `HF_HUB_OFFLINE=1`. Do not omit the audio tokenizer:
+the main model directory does not contain its weights.
+
+## Start the server
+
+Expose one physical NPU and keep the stage's logical device as `0`:
+
+```bash
+export ASCEND_RT_VISIBLE_DEVICES=0
+export HF_HOME=/models/huggingface-cache
+
+vllm serve OpenMOSS-Team/MOSS-TTS-Nano \
+    --omni \
+    --trust-remote-code \
+    --host 0.0.0.0 \
+    --port 8091
+```
+
+`--trust-remote-code` is required because both checkpoints provide custom
+Transformers model implementations. The bundled `moss_tts_nano.yaml` deploy
+configuration enables eager execution; ACL graph capture is not used for this
+model path.
+
+## Voice-cloning request
+
+`ref_audio` accepts an HTTP URL or a base64 data URL. A non-streaming request
+returns a WAV file:
+
+```bash
+curl http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "OpenMOSS-Team/MOSS-TTS-Nano",
+        "input": "Hello, this is a voice cloning test on Ascend NPU.",
+        "stream": false,
+        "response_format": "wav",
+        "ref_audio": "https://example.com/reference.wav"
+    }' \
+    --output output.wav
+```
+
+Streaming raw PCM uses the same endpoint:
+
+```bash
+curl http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "OpenMOSS-Team/MOSS-TTS-Nano",
+        "input": "你好，这是一段流式语音合成测试。",
+        "stream": true,
+        "stream_format": "audio",
+        "response_format": "pcm",
+        "ref_audio": "https://example.com/reference.wav"
+    }' \
+    --output output.pcm
+```
+
+Raw PCM output is signed 16-bit little-endian mono at 48 kHz.
+
+## Validation
+
+The Ascend 910B3 validation covered:
+
+- English non-streaming WAV output;
+- Chinese streaming PCM output;
+- repeat generation with the same seed producing byte-identical WAV files;
+- two concurrent non-streaming requests;
+- waveform checks for a non-empty, finite, non-silent signal at 48 kHz.
+
+The first request includes one-time NPU kernel warmup. In the validated setup,
+the first non-streaming request took about 22.7 seconds, a repeated request
+took about 7.5 seconds, and streaming produced its first audio chunk in about
+0.38 seconds. These figures are functional reference measurements, not a
+performance guarantee.
+
+## Current limitations
+
+- The model runs in eager mode; ACL graph mode has not been enabled.
+- Concurrent HTTP requests are accepted but generated serially. Do not raise
+  the stage's `max_num_seqs` above `1`: the remote model and audio tokenizer
+  share RNG and streaming decode state, and interleaved generation can corrupt
+  audio.
+- The remote model's CUDA-only `flash_attention_2` path is replaced with
+  PyTorch SDPA on Ascend.
+- The current non-CUDA model wrapper loads the language model and audio codec
+  in float32. Quantized checkpoints have not been validated.
+- A reference audio clip is required for `voice_clone` mode.
@@ -74,6 +74,7 @@ recipes/
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
 | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
 | [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
+| [`OpenMOSS/MOSS-TTS-Nano-NPU.md`](./OpenMOSS/MOSS-TTS-Nano-NPU.md) | Online multilingual voice-cloned TTS | 1x Ascend 910B3 64GB |
 | [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
 | [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
 | [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk |
@@ -121,7 +121,7 @@ def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:


@pytest.mark.advanced_model
-@hardware_test(res={"cuda": "L4"})
+@hardware_test(res={"cuda": "L4", "npu": "A2"})
 def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None:
    """English TTS produces non-empty 48 kHz stereo audio."""
    req = _build_request("Hello, this is a short voice cloning demo for testing.", ref_audio_path)
@@ -133,7 +133,7 @@ def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None:


@pytest.mark.advanced_model
-@hardware_test(res={"cuda": "L4"})
+@hardware_test(res={"cuda": "L4", "npu": "A2"})
 def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None:
    """Chinese TTS produces non-empty audio."""
    req = _build_request("你好，这是语音合成测试。", ref_audio_path)
@@ -145,7 +145,7 @@ def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None:


@pytest.mark.advanced_model
-@hardware_test(res={"cuda": "L4"})
+@hardware_test(res={"cuda": "L4", "npu": "A2"})
 def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) -> None:
    """Same seed produces identical waveforms."""
    req = _build_request("Reproducible output test.", ref_audio_path, seed=123)
@@ -157,7 +157,7 @@ def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) ->


@pytest.mark.advanced_model
-@hardware_test(res={"cuda": "L4"})
+@hardware_test(res={"cuda": "L4", "npu": "A2"})
 def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None:
    """Batch of two requests returns audio for each."""
    requests = [
@@ -84,7 +84,7 @@ tts_server_params = [
 ]


-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "npu": "A2"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> None:
    """
@@ -106,7 +106,7 @@ def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> No
    online_client.send_audio_speech_request(request_config)


-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "npu": "A2"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_text_to_audio_002(omni_server, online_client, ref_audio_data_url) -> None:
    """
@@ -138,7 +138,7 @@ def test_text_to_audio_002(omni_server, online_client, ref_audio_data_url) -> No
    online_client.send_audio_speech_request(request_config)


-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "npu": "A2"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_text_to_audio_003(omni_server, online_client, ref_audio_data_url) -> None:
    """
@@ -160,7 +160,7 @@ def test_text_to_audio_003(omni_server, online_client, ref_audio_data_url) -> No
    online_client.send_audio_speech_request(request_config)


-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "npu": "A2"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_text_to_audio_004_ref_text_ignored(omni_server, online_client, ref_audio_data_url) -> None:
    """
@@ -181,3 +181,23 @@ def test_text_to_audio_004_ref_text_ignored(omni_server, online_client, ref_audi
    }

    online_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "L4", "npu": "A2"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_005_concurrent_requests_are_serialized(omni_server, online_client, ref_audio_data_url) -> None:
+    """Concurrent clients produce identical audio when generation is serialized."""
+    request_config = {
+        "model": omni_server.model,
+        "input": "Concurrent clients should receive clean and deterministic speech.",
+        "stream": False,
+        "response_format": "wav",
+        "ref_audio": ref_audio_data_url,
+        "seed": 42,
+    }
+
+    responses = online_client.send_audio_speech_request(request_config, request_num=2)
+
+    assert len(responses) == 2
+    assert responses[0].audio_bytes
+    assert responses[0].audio_bytes == responses[1].audio_bytes
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts_nano.modeling_moss_tts_nano import (
+    _configure_attention_implementation,
+    _validate_max_num_seqs,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_non_cuda_attention_uses_sdpa():
+    class FakeModel:
+        attention_implementation = None
+
+        def _set_attention_implementation(self, implementation):
+            self.attention_implementation = implementation
+
+    model = FakeModel()
+    _configure_attention_implementation(model, torch.device("cpu"))
+
+    assert model.attention_implementation == "sdpa"
+
+
+def test_moss_tts_nano_rejects_concurrent_active_sequences():
+    _validate_max_num_seqs(1)
+
+    with pytest.raises(ValueError, match="requires max_num_seqs=1"):
+        _validate_max_num_seqs(2)
@@ -6,8 +6,9 @@
 #   * enforce_eager=true — the model uses upstream inference_stream() with
 #     trust_remote_code; CUDA graph capture is not wired up.
 #   * gpu_memory_utilization=0.3 — the AR LM + codec fit in ~2 GiB.
-#   * max_num_seqs=4 — supports per-request streaming generators stored in
-#     self._stream_gens.
+#   * max_num_seqs=1 — the remote model and audio tokenizer keep process-wide
+#     RNG and streaming decode state. HTTP requests may arrive concurrently,
+#     but generation must remain serial to avoid cross-request audio corruption.
 #
 # MOSS-TTS-Nano has no multimodal input preprocessing, so skip_mm_profiling
 # bypasses the dummy-MM pass at profile time.
@@ -16,7 +17,7 @@ trust_remote_code: true

 stages:
  - stage_id: 0
-    max_num_seqs: 4
+    max_num_seqs: 1
    gpu_memory_utilization: 0.3
    enforce_eager: true
    enable_prefix_caching: false
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """MOSS-TTS-Nano single-stage model for vLLM-Omni.

 Runs in a single AR worker stage.  The 0.1B AR LM and the
@@ -160,6 +160,33 @@ def _pick(info: dict, key: str, default):
    return val if val is not None else default


+def _configure_attention_implementation(lm: nn.Module, device: torch.device) -> None:
+    if device.type == "cuda":
+        try:
+            import flash_attn  # noqa: F401
+
+            lm._set_attention_implementation("flash_attention_2")
+            logger.info("MOSS-TTS-Nano using flash_attention_2")
+        except ImportError:
+            lm._set_attention_implementation("sdpa")
+            logger.info("MOSS-TTS-Nano using sdpa (flash_attn not installed)")
+    else:
+        # The checkpoint defaults to flash_attention_2, whose external
+        # flash_attn package is CUDA-only. PyTorch SDPA is dispatched to the
+        # native backend on non-CUDA accelerators such as Ascend NPU.
+        lm._set_attention_implementation("sdpa")
+        logger.info("MOSS-TTS-Nano using sdpa on %s", device.type)
+
+
+def _validate_max_num_seqs(max_num_seqs: int) -> None:
+    if max_num_seqs != 1:
+        raise ValueError(
+            "MOSS-TTS-Nano currently requires max_num_seqs=1 because its "
+            "remote model and audio tokenizer share RNG and streaming decode "
+            "state across requests. Higher values can corrupt generated audio."
+        )
+
+
 class MossTTSNanoForGeneration(nn.Module):
    """Single-stage MOSS-TTS-Nano model with streaming audio output.

@@ -179,6 +206,7 @@ class MossTTSNanoForGeneration(nn.Module):
        self.vllm_config = vllm_config
        self.config = vllm_config.model_config.hf_config
        self.model_path: str = vllm_config.model_config.model
+        _validate_max_num_seqs(vllm_config.scheduler_config.max_num_seqs)

        # Eager construction (not in load_weights) -- see module docstring.
        _patch_torchaudio_load()
@@ -209,15 +237,7 @@ class MossTTSNanoForGeneration(nn.Module):
                trust_remote_code=True,
                torch_dtype=tts_dtype,
            )
-        if device.type == "cuda":
-            try:
-                import flash_attn  # noqa: F401
-
-                lm._set_attention_implementation("flash_attention_2")
-                logger.info("MOSS-TTS-Nano using flash_attention_2")
-            except ImportError:
-                lm._set_attention_implementation("sdpa")
-                logger.info("MOSS-TTS-Nano using sdpa (flash_attn not installed)")
+        _configure_attention_implementation(lm, device)
        lm.to(device=device)
        lm.eval()