[Hardware][Ascend][Model] Add NPU support for MiniCPM-o 4.5 talker/vo… (#5117)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

examples/online_serving/minicpmo/README.md

@@ -11,17 +11,18 @@ The UI supports:
 ## 1. Start the backend server
 
 The deploy config auto-loads via `--omni`; the default
-`vllm_omni/deploy/minicpmo_4_5.yaml` targets a 2-GPU layout (thinker on
-GPU 0, talker + t2w sharing GPU 1).  For other hardware layouts pick
+`vllm_omni/deploy/minicpmo_4_5.yaml` targets a single-GPU layout (thinker
+and talker + t2w co-located on GPU 0).  For other hardware layouts pick
 one of the deploy variants below.
 
 | deploy config | GPUs | Notes |
 |---|---|---|
-| `minicpmo_4_5.yaml` (default) | 2 | Thinker on GPU0, talker+t2w on GPU1. |
+| `minicpmo_4_5.yaml` (default) | 1 | Thinker and talker+t2w co-located on GPU0. |
+| `minicpmo_4_5_2gpu.yaml` | 2 | Thinker on GPU0, talker+t2w on GPU1. |
 | `minicpmo_4_5_3gpu.yaml` | 3 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. |
 | `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
 
-Default (2-GPU):
+Default (single-GPU):
 
 ```bash
 vllm serve openbmb/MiniCPM-o-4_5 --omni \

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

@@ -19,15 +19,18 @@ Use this recipe as a known-good starting point for serving
 of the MiniCPM-o family — it pairs a multimodal-understanding thinker
 LLM with a streaming `MiniCPMTTS + Token2Wav` talker so a single
 `/v1/chat/completions` call can return text and 24 kHz speech in one
-shot. The recipe covers three shipped GPU layouts (2 / 3 / 8 GPUs)
-selected via `--deploy-config`.
+shot. The recipe covers the shipped GPU layouts (single / 2 / 3 / 8 GPUs):
+  the default co-locates both stages on one GPU, and the larger scale-out
+  layouts (2 / 3 / 8 GPUs) are selected via `--deploy-config`.
 
 ## References
 
 - Default deploy configs (auto-loaded by HF `model_type=minicpmo` +
   `hf_config.version="4.5"`):
-  - 2-GPU layout (default):
+  - Single-GPU layout (default):
     [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
+  - 2-GPU layout (thinker on GPU 0, talker on GPU 1):
+    [`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml)
   - 3-GPU layout (thinker TP=2):
     [`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
   - 8x RTX 4090 layout:
@@ -45,28 +48,30 @@ selected via `--deploy-config`.
 
 ## Hardware Support
 
-Three GPU layouts ship with default deploy configs. Pick the layout that
-matches your hardware and pass it via `--deploy-config`; the talker
-(`MiniCPMTTS + Token2Wav`) always lives on its own GPU because of the
-in-process vocoder, and the thinker is the part that scales out via TP.
+Four GPU layouts ship with default deploy configs. Pick the layout that
+matches your hardware and pass it via `--deploy-config`. The default
+co-locates the thinker and the talker (`MiniCPMTTS + Token2Wav`, with its
+in-process vocoder) on a single GPU; the multi-GPU layouts split the
+talker onto its own GPU and scale the thinker out via TP.
 
 | Layout | Thinker | Talker + Token2Wav | Typical hardware |
 | --- | --- | --- | --- |
-| 2-GPU (default) | GPU 0 | GPU 1 | 2x A100/H100/H200 80GB |
+| Single-GPU (default) | GPU 0 | GPU 0 | 1x A100/H100/H200 80GB |
+| 2-GPU | GPU 0 | GPU 1 | 2x A100/H100/H200 80GB |
 | 3-GPU (thinker TP=2) | GPU 0,1 (TP=2) | GPU 2 | 3x mid-tier GPUs |
 | 8x RTX 4090 24GB | GPU 0–3 (TP=4) | GPU 4 | 8x RTX 4090 consumer |
 
 ## GPU
 
-### 2 x GPU (default — single command)
+### 1 x GPU (default — single command)
 
 The default
 [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
-puts the thinker on GPU 0 (`~70 %` memory, `enforce_eager: true`,
-`max_num_seqs: 1`) and the talker + Token2Wav vocoder on GPU 1
-(`~75 %` memory). This is the recommended starting layout — works on
-any pair of 80GB-class GPUs (A100, H100, H200) and on most 40GB+
-pairs as long as the thinker model weights fit.
+co-locates both stages on GPU 0: the thinker (`~80 %` memory) and the
+talker + Token2Wav vocoder (`~10 %` memory, `max_num_seqs: 1`). This is
+the recommended starting layout — works on a single 80GB-class GPU
+(A100, H100, H200) as long as the thinker model weights and the talker's
+in-process vocoder fit together.
 
 #### Environment
 
@@ -85,7 +90,7 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
 ```
 
 The deploy config is auto-loaded by the model registry — no
-`--deploy-config` flag needed for this default 2-GPU layout.
+`--deploy-config` flag needed for this default single-GPU layout.
 
 #### Verification
 
@@ -142,17 +147,40 @@ speech output (TTS)"** checkbox on / off.
 
 #### Notes
 
-- Memory budget: thinker weights occupy GPU 0 at `gpu_memory_utilization:
-  0.7`; talker + Token2Wav vocoder share GPU 1 at `0.75`.
+- Memory budget: both stages share GPU 0 — the thinker at
+  `gpu_memory_utilization: 0.8`, the talker + Token2Wav vocoder at `0.1`.
 - `--trust-remote-code` is required — the HF repo ships a custom
   `MiniCPMO` config / model class.
-- Pin: `enforce_eager: true` on both stages (CUDA graph capture is off
-  by design for the talker's Token2Wav path).
 - Stage 1 (talker) is hard-capped to `max_num_seqs: 1`: the talker
   only consumes `runtime_additional_information[0]`, so any value > 1
   makes concurrent requests share request-0's audio. This is the same
   cap baked into the deploy config.
 
+### 2 x GPU (talker on its own GPU)
+
+Use
+[`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml)
+when you have two GPUs and want to give the talker + Token2Wav vocoder a
+dedicated card instead of sharing GPU 0 with the thinker. The thinker
+runs on GPU 0 (`~90 %` mem, TP=1) and the talker on GPU 1 (`~75 %` mem,
+`max_num_seqs: 1`). This relieves the memory pressure of the default
+single-GPU co-located layout and is the recommended step up when a
+second 80GB-class card is available but full 3-way TP scale-out is not
+needed.
+
+#### Command
+
+```bash
+vllm serve openbmb/MiniCPM-o-4_5 --omni \
+    --deploy-config vllm_omni/deploy/minicpmo_4_5_2gpu.yaml \
+    --trust-remote-code \
+    --host 0.0.0.0 --port 8099
+```
+
+Verification and Notes mirror the single-GPU section; the only
+difference is that the talker no longer competes with the thinker for
+GPU 0 memory.
+
 ### 3 x GPU (thinker TP=2)
 
 Use
@@ -171,7 +199,7 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
     --host 0.0.0.0 --port 8099
 ```
 
-Verification and Notes mirror the 2-GPU section; thinker latency
+Verification and Notes mirror the single-GPU section; thinker latency
 roughly halves under load thanks to TP=2.
 
 ### 8 x RTX 4090 24GB (consumer-GPU layout)
@@ -197,8 +225,8 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
   4090s. Raise it if your cards have more headroom (e.g. 4090 D /
   custom 32 GB SKUs), but verify with a long-prompt run before
   promoting.
-- All other knobs match the 2-GPU section; the only difference is the
-  per-card memory pressure on the thinker shards.
+- All other knobs match the single-GPU section; the only difference is
+  the per-card memory pressure on the thinker shards.
 
 ## Notes (applies to all layouts)
 

vllm_omni/deploy/minicpmo_4_5.yaml

@@ -1,9 +1,9 @@
-# MiniCPM-o 4.5 default deploy: 2-GPU layout.
+# MiniCPM-o 4.5 default deploy: single-GPU layout.
 #
-# Stage 0 (thinker): GPU 0 — multimodal understanding, ~70% mem.
-# Stage 1 (talker) : GPU 1 — MiniCPMTTS + Token2Wav, ~75% mem
-#                            (covers both the talker LM and the in-process
-#                            Token2Wav vocoder that emits the waveform).
+# Both stages are co-located on GPU 0:
+#   Stage 0 (thinker): multimodal understanding, ~80% mem.
+#   Stage 1 (talker) : MiniCPMTTS + in-process Token2Wav vocoder, ~10% mem
+#                      (talker LM plus the vocoder that emits the waveform).
 #
 # Other GPU layouts: minicpmo_4_5_3gpu.yaml, minicpmo_4_5_8x4090.yaml.
 async_chunk: false
@@ -13,9 +13,9 @@ enable_prefix_caching: false
 
 stages:
   - stage_id: 0
-    max_num_seqs: 1
-    gpu_memory_utilization: 0.7
-    enforce_eager: true
+    max_num_seqs: 16
+    gpu_memory_utilization: 0.8
+    enforce_eager: false
     tensor_parallel_size: 1
     max_num_batched_tokens: 16384
     devices: "0"
@@ -34,10 +34,10 @@ stages:
     # Talker only consumes runtime_additional_information[0]; any value > 1
     # makes concurrent requests share request-0's audio (mm_outputs fallback).
     max_num_seqs: 1
-    gpu_memory_utilization: 0.75
-    enforce_eager: true
+    gpu_memory_utilization: 0.1
+    enforce_eager: false
     max_num_batched_tokens: 8192
-    devices: "1"
+    devices: "0"
     default_sampling_params:
       temperature: 0.0
       top_p: 1.0
@@ -46,6 +46,14 @@ stages:
       seed: 42
       detokenize: false
 
-# MiniCPM-o 4.5 has only been verified on CUDA devices.  Add per-platform
-# overrides here when other backends (ROCm/XPU/NPU) are verified end-to-end.
-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

vllm_omni/deploy/minicpmo_4_5_2gpu.yaml

@@ -0,0 +1,64 @@
+# MiniCPM-o 4.5 deploy: 2-GPU layout (thinker on GPU 0, talker on GPU 1).
+#
+# Stage 0 (thinker): GPU 0 — multimodal understanding, ~90% mem (thinker
+#                            gets a dedicated card, so it can take more of
+#                            the GPU than the single-GPU co-located layout).
+# Stage 1 (talker) : GPU 1 — MiniCPMTTS + Token2Wav, ~75% mem
+#                            (covers both the talker LM and the in-process
+#                            Token2Wav vocoder that emits the waveform).
+#
+# Use this when you have two GPUs and want the talker + Token2Wav vocoder
+# on its own card instead of sharing GPU 0 with the thinker (the default
+# single-GPU layout co-locates both and can OOM under load).
+pipeline: minicpmo_4_5
+async_chunk: false
+
+trust_remote_code: true
+enable_prefix_caching: false
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 16
+    gpu_memory_utilization: 0.9
+    enforce_eager: false
+    tensor_parallel_size: 1
+    max_num_batched_tokens: 16384
+    devices: "0"
+    limit_mm_per_prompt:
+      video: 1
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.1
+
+  - stage_id: 1
+    # Talker only consumes runtime_additional_information[0]; any value > 1
+    # makes concurrent requests share request-0's audio (mm_outputs fallback).
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.75
+    enforce_eager: false
+    max_num_batched_tokens: 8192
+    devices: "1"
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 1
+      seed: 42
+      detokenize: false
+
+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

vllm_omni/deploy/minicpmo_4_5_3gpu.yaml

@@ -12,9 +12,9 @@ enable_prefix_caching: false
 
 stages:
   - stage_id: 0
-    max_num_seqs: 1
+    max_num_seqs: 16
     gpu_memory_utilization: 0.7
-    enforce_eager: true
+    enforce_eager: false
     tensor_parallel_size: 2
     max_num_batched_tokens: 16384
     devices: "0,1"
@@ -34,7 +34,7 @@ stages:
     # makes concurrent requests share request-0's audio (mm_outputs fallback).
     max_num_seqs: 1
     gpu_memory_utilization: 0.75
-    enforce_eager: true
+    enforce_eager: false
     max_num_batched_tokens: 8192
     devices: "2"
     default_sampling_params:

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -760,7 +760,12 @@ class MiniCPMVImageProcessor(BaseImageProcessor):
         )
 
 
-AutoImageProcessor.register("MiniCPMVImageProcessor", MiniCPMVImageProcessor)
+# transformers >= 5.x requires the config *class* (not a string) as the
+# first arg to AutoImageProcessor.register — it does `key.__module__`
+# internally. Passing "MiniCPMVImageProcessor" raises:
+#   AttributeError: 'str' object has no attribute '__module__'
+# See also processing_gr00t_n1d7.py for the same transformers API change.
+AutoImageProcessor.register(MiniCPMOConfig, MiniCPMVImageProcessor, exist_ok=True)
 
 
 # ============== SigLIP Vision Transformer Classes ==============

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py

@@ -15,6 +15,7 @@ Pipeline:
 import io
 import logging
 import os
+import sys
 from collections.abc import Iterable
 
 import numpy as np
@@ -27,13 +28,33 @@ from vllm.model_executor.models.interfaces import SupportsPP
 
 from vllm_omni.platforms import current_omni_platform
 
-try:
-    from stepaudio2 import Token2wav as _Token2wav
+# Preserve the established external vocoder on CUDA. Ascend uses the in-tree
+# adapter because ``stepaudio2-minicpmo`` hard-codes CUDA device placement.
+if current_omni_platform.is_npu():
+    try:
+        from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_token2wav import (
+            MiniCPMO45Token2wav as _Token2wav,
+        )
 
-    _stepaudio2_available = True
-except ImportError:
-    _Token2wav = None
-    _stepaudio2_available = False
+        _token2wav_backend = "step_audio2_core"
+    except ImportError:
+        try:
+            from stepaudio2 import Token2wav as _Token2wav
+
+            _token2wav_backend = "stepaudio2_pkg"
+        except ImportError:
+            _Token2wav = None
+            _token2wav_backend = None
+else:
+    try:
+        from stepaudio2 import Token2wav as _Token2wav
+
+        _token2wav_backend = "stepaudio2_pkg"
+    except ImportError:
+        _Token2wav = None
+        _token2wav_backend = None
+
+_stepaudio2_available = _Token2wav is not None
 
 logger = logging.getLogger(__name__)
 
@@ -101,14 +122,25 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
             return
         try:
             model_path = self.vllm_config.model_config.model
-            import os
-            import sys
 
             if model_path not in sys.path:
                 sys.path.insert(0, model_path)
+            from transformers import AutoImageProcessor
             from transformers.dynamic_module_utils import get_class_from_dynamic_module
 
-            MiniCPMTTS = get_class_from_dynamic_module("modeling_minicpmo.MiniCPMTTS", model_path)
+            # openbmb/MiniCPM-o-4_5/processing_minicpmo.py registers via a
+            # string: AutoImageProcessor.register("MiniCPMVImageProcessor", ...),
+            # which crashes on transformers>=5 (register reads key.__module__).
+            # Loading MiniCPMTTS imports that module, so no-op the string form
+            # (unused by the standalone talker) while it runs, then restore.
+            original_register = AutoImageProcessor.register
+            AutoImageProcessor.register = (  # type: ignore[method-assign]
+                lambda key, *a, **k: None if isinstance(key, str) else original_register(key, *a, **k)
+            )
+            try:
+                MiniCPMTTS = get_class_from_dynamic_module("modeling_minicpmo.MiniCPMTTS", model_path)
+            finally:
+                AutoImageProcessor.register = original_register  # type: ignore[method-assign]
 
             # MiniCPMTTS.__init__ reads `config.top_p / top_k / repetition_penalty`
             # directly (modeling_minicpmo.py L4112-4114), but the model repo's
@@ -121,6 +153,11 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
                 if not hasattr(self._tts_config, _attr):
                     setattr(self._tts_config, _attr, _default)
 
+            # The copied Hugging Face flash_attention_2 setting is not valid
+            # for this standalone MiniCPMTTS path. Use PyTorch SDPA on every
+            # backend until a dedicated flash-attention implementation exists.
+            self._tts_config.attn_implementation = "sdpa"
+
             prev_dtype = torch.get_default_dtype()
             torch.set_default_dtype(torch.float32)
             try:
@@ -159,7 +196,11 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
                 finally:
                     torch.set_default_dtype(prev_dtype2)
                 self.tts_obj.audio_tokenizer = self.audio_tokenizer
-                logger.info("Loaded Token2wav from %s", token2wav_dir)
+                logger.info(
+                    "Loaded Token2wav from %s (backend=%s)",
+                    token2wav_dir,
+                    _token2wav_backend,
+                )
             # Only mark init as complete after every step succeeds, so a
             # partial failure leaves the next call free to retry the full
             # init instead of short-circuiting back to a silent empty path.
@@ -189,19 +230,25 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
 
         tts = self.tts_obj
         device = tts.emb_text.weight.device
-        dtype = tts.emb_text.weight.dtype
+        # MiniCPMTTS AR backbone uses FlashAttention (fp16/bf16 only). The
+        # submodule is constructed under float32 default dtype during lazy init,
+        # so pin the condition embeddings to bfloat16 explicitly rather than
+        # inheriting the (float32) parameter dtype — a float32 condition breaks
+        # the CUDA FA2 path and wastes memory on the NPU sdpa path.
+        ar_dtype = torch.bfloat16
 
         llm_embeds = tts.emb_text(tts_token_ids.to(device))
-        hidden_embeds = tts.projector_semantic(tts_hidden_states.to(device=device, dtype=dtype))
+        hidden_embeds = tts.projector_semantic(tts_hidden_states.to(device=device, dtype=ar_dtype))
         if getattr(tts.config, "normalize_projected_hidden", False):
             hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1)
-        tts_embeds = llm_embeds + hidden_embeds
+        tts_embeds = (llm_embeds + hidden_embeds).to(dtype=ar_dtype)
 
         text_eos = tts.emb_text(torch.tensor([tts.config.text_eos_token_id], device=device, dtype=torch.long))
         audio_bos = tts.emb_text(torch.tensor([tts.audio_bos_token_id], device=device, dtype=torch.long))
-        spk_embeds = torch.zeros(0, tts.config.hidden_size, device=device, dtype=tts_embeds.dtype)
+        spk_embeds = torch.zeros(0, tts.config.hidden_size, device=device, dtype=ar_dtype)
 
         inputs_embeds = torch.cat([spk_embeds, tts_embeds, text_eos, audio_bos], dim=0).unsqueeze(0)
+        inputs_embeds = inputs_embeds.to(dtype=ar_dtype)
         logger.info("generate_speech: inputs_embeds shape=%s", list(inputs_embeds.shape))
 
         # Scale max_new_token with input text length to avoid mid-stream truncation on long
@@ -249,7 +296,15 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
         prev_dtype = torch.get_default_dtype()
         torch.set_default_dtype(torch.float32)
         try:
-            with torch.amp.autocast("cuda", enabled=False):
+            # Vocoder path is float32; use the platform abstraction because
+            # torch.amp.autocast validates unsupported device types even when
+            # autocast is disabled.
+            autocast_device = device.type if isinstance(device, torch.device) else str(device)
+            with current_omni_platform.create_autocast_context(
+                device_type=autocast_device,
+                dtype=torch.float32,
+                enabled=False,
+            ):
                 token_list = generated_tokens.squeeze(0).tolist()
                 num_tokens = len(token_list)
 
@@ -418,7 +473,17 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
         return None, None
 
     def compute_logits(self, hidden_states, *args, **kwargs):
-        return torch.zeros(1, 2, device=hidden_states.device if isinstance(hidden_states, torch.Tensor) else "cuda")
+        # Placeholder logits: one row per sampled request (the scheduler
+        # indexes sampled_token_ids by req_index). Hardcoding a single row
+        # breaks batched/concurrent decoding with IndexError. The values are
+        # discarded — real output is the waveform via multimodal_outputs.
+        if isinstance(hidden_states, torch.Tensor):
+            device = hidden_states.device
+            num_reqs = hidden_states.shape[0] if hidden_states.ndim >= 1 else 1
+        else:
+            device = current_omni_platform.get_torch_device()
+            num_reqs = 1
+        return torch.zeros(num_reqs, 2, device=device)
 
     def sample(self, logits, sampling_metadata):
         return None
@@ -442,10 +507,31 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
                     logger.warning("TTS missing keys (%d): %s", len(missing), missing[:5])
                 if unexpected:
                     logger.warning("TTS unexpected keys (%d): %s", len(unexpected), unexpected[:5])
-                self.tts_obj = self.tts_obj.to("cuda")
+                # Move the AR backbone to the active device (cuda / npu / …) and
+                # cast to bfloat16: MiniCPMTTS AR uses FlashAttention (fp16/bf16
+                # only) and is built under a float32 default dtype during lazy
+                # init, so an uncast float32 backbone breaks CUDA FA2 and wastes
+                # memory on the NPU sdpa path. Detach the Token2wav vocoder first
+                # so the cast does not drag it onto the accelerator or downcast
+                # its float32 flow/HiFT weights: it manages its own device
+                # placement (e.g. CPU HiFT on NPU) and may not be an nn.Module.
+                device = current_omni_platform.get_torch_device()
+                audio_tok = getattr(self.tts_obj, "audio_tokenizer", None)
+                if audio_tok is not None:
+                    self.tts_obj.audio_tokenizer = None
+                try:
+                    self.tts_obj = self.tts_obj.to(device=device, dtype=torch.bfloat16)
+                finally:
+                    if audio_tok is not None:
+                        self.tts_obj.audio_tokenizer = audio_tok
+                        self.audio_tokenizer = audio_tok
                 self.emb_text = self.tts_obj.emb_text
                 self.projector_semantic = self.tts_obj.projector_semantic
-                logger.info("Loaded %d TTS weights, moved to cuda", len(tts_weights))
+                logger.info(
+                    "Loaded %d TTS weights, moved to %s (bfloat16)",
+                    len(tts_weights),
+                    device,
+                )
 
         return loaded
 

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniCPM-o 4.5 Token2wav adapter over in-tree ``StepAudio2Token2WavCore``.
+
+``minicpmo_4_5_omni_tts`` historically depended on the external
+``from stepaudio2 import Token2wav`` entry point (``stepaudio2-minicpmo``).
+That package hard-codes ``.cuda()`` and duplicates the flow/HiFT stack that
+vLLM-Omni already vendors for Step-Audio2 (which also carries the Ascend/NPU
+fixes: CPU HiFT, DiT mask expand, MATH SDPA, compile disable, PE buffer
+extension).
+
+This module exposes the same call surface MiniCPM expects
+(``__call__`` / ``set_stream_cache`` / ``stream``) while delegating the
+actual vocoder work to
+``vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav.StepAudio2Token2WavCore``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+
+from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import (
+    StepAudio2Token2WavCore,
+    _StreamState,
+)
+
+
+def _resolve_device(device: str | torch.device | None) -> str:
+    if device is not None:
+        return str(device)
+    try:
+        from vllm_omni.platforms import current_omni_platform
+
+        plat_dev = getattr(current_omni_platform, "device_type", None)
+        if plat_dev:
+            return str(plat_dev)
+    except Exception:
+        pass
+    if torch.cuda.is_available():
+        return "cuda"
+    if hasattr(torch, "npu") and torch.npu.is_available():
+        return "npu"
+    return "cpu"
+
+
+class MiniCPMO45Token2wav:
+    """Token2wav-compatible facade around ``StepAudio2Token2WavCore``.
+
+    Matches the MiniCPM-o / ``stepaudio2-minicpmo`` vocoder API used by
+    ``MiniCPMO45OmniTTSForConditionalGeneration.generate_speech``:
+
+    * ``Token2wav(model_path, float16=False, n_timesteps=10)``
+    * ``__(tokens, prompt_wav) -> wav_bytes``
+    * ``set_stream_cache(prompt_wav) -> (stream_cache, hift_cache_dict)``
+    * ``stream(tokens, prompt_wav, last_chunk=..., return_waveform=...)``
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        float16: bool = False,
+        n_timesteps: int = 10,
+        device: str | torch.device | None = None,
+    ):
+        self.float16 = float16
+        self.n_timesteps = n_timesteps
+        self.device = _resolve_device(device)
+        self._core = StepAudio2Token2WavCore(
+            model_path=model_path,
+            float16=float16,
+            device=self.device,
+            n_timesteps=n_timesteps,
+        )
+        # Eager-load so construction failures surface at init time (same as
+        # the external Token2wav package), not on the first request.
+        self._core._ensure_models_loaded()
+
+        # Mutable streaming fields expected by MiniCPM's long-form path.
+        self.stream_cache: Any | None = None
+        self.hift_cache_dict: dict[str, torch.Tensor] = {}
+
+    def __call__(self, generated_speech_tokens, prompt_wav) -> bytes:
+        """One-shot tokens → 24 kHz WAV bytes."""
+        return self._core.forward(
+            generated_speech_tokens,
+            prompt_wav,
+            return_bytes=True,
+        )
+
+    def set_stream_cache(self, prompt_wav: str):
+        """Initialise flow + HiFT caches for chunked vocoding.
+
+        Returns ``(stream_cache, hift_cache_dict)`` so callers can assign them
+        back onto this object the way ``stepaudio2.Token2wav`` does.
+        """
+        state = _StreamState()
+        self._core.setup_stream_for(prompt_wav, state)
+        self.stream_cache = state.stream_cache
+        self.hift_cache_dict = state.hift_cache_dict
+        return self.stream_cache, self.hift_cache_dict
+
+    def stream(
+        self,
+        generated_speech_tokens,
+        prompt_wav: str,
+        last_chunk: bool = False,
+        return_waveform: bool = False,
+    ):
+        """Process one streaming chunk; updates ``stream_cache`` / ``hift_cache_dict``."""
+        if self.stream_cache is None:
+            raise ValueError("stream_cache is not set")
+
+        state = _StreamState()
+        state.setup_done = True
+        state.stream_cache = self.stream_cache
+        state.hift_cache_dict = self.hift_cache_dict
+
+        speech = self._core.stream_chunk_for(
+            list(generated_speech_tokens),
+            prompt_wav,
+            last_chunk,
+            state,
+        )
+        self.stream_cache = state.stream_cache
+        self.hift_cache_dict = state.hift_cache_dict
+
+        wav_np = speech.detach().float().cpu().numpy()
+        if return_waveform:
+            return wav_np
+
+        wav_np = np.clip(wav_np, -1.0, 1.0)
+        wav_int16 = (wav_np * 32767.0).astype("<i2")
+        return wav_int16.tobytes()

vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py

@@ -95,6 +95,22 @@ class StepAudio2Token2WavCore(nn.Module):
         self.source_cache_len = STREAM_SOURCE_CACHE_LEN
         self.speech_window: torch.Tensor | None = None  # created lazily on device
 
+        # On Ascend the talker worker never routes through
+        # ``NPUOmniPlatform.set_device``, so gate on ``current_omni_platform``
+        # and apply the Token2Wav NPU patch here (idempotent) before
+        # ``_ensure_models_loaded`` runs. Covers every
+        # construction site (Step-Audio2 wrapper + MiniCPM-o facade): patched
+        # ``_ensure_models_loaded`` offloads HiFT to CPU and patched ``forward``
+        # expands CosyVoice DiT masks + forces MATH SDPA (avoids FA 161001).
+        from vllm_omni.platforms import current_omni_platform
+
+        if current_omni_platform.is_npu():
+            from vllm_omni.platforms.npu.models.step_audio2_token2wav import (
+                apply_step_audio2_token2wav_npu_patch,
+            )
+
+            apply_step_audio2_token2wav_npu_patch()
+
     def _ensure_models_loaded(self):
         """Lazy load models on first use"""
         if self._models_loaded:

vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NPU patches for CosyVoice2 / Token2Wav DiT attention.
+
+CosyVoice2's DiT ``Attention.forward`` builds a key-padding style mask
+``(B, 1, 1, S)`` via ``mask.unsqueeze(1)`` and feeds it to
+``F.scaled_dot_product_attention``. On Ascend that call is routed to
+``npu_fusion_attention`` / ``aclnnFlashAttentionScore``, which only accepts
+mask shapes ``[B,N,S,S] / [B,1,S,S] / [1,1,S,S] / [S,S]`` and fails with
+error **161001** (tiling / parameter invalid) for ``[B,1,1,S]``.
+
+This module:
+1. Expands DiT attention masks to ``[B, 1, S, S]`` before SDPA.
+2. Provides a MATH-backend SDPA context so inference can avoid the fused FA
+   kernel entirely when the platform still incorrectly routes SDPA.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager, nullcontext
+
+import torch
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_PATCHED = False
+
+
+def _expand_attn_mask_for_npu(
+    attn_mask: torch.Tensor | None,
+    q_len: int,
+    kv_len: int | None = None,
+) -> torch.Tensor | None:
+    """Expand CosyVoice key-padding masks to Ascend FA-compatible shapes."""
+    if attn_mask is None:
+        return None
+    kv_len = kv_len if kv_len is not None else q_len
+
+    # (B, 1, S) key-padding -> (B, S_q, S_kv) then unsqueeze heads below.
+    if attn_mask.dim() == 3 and attn_mask.shape[-2] == 1:
+        attn_mask = attn_mask.expand(-1, q_len, -1)
+    if attn_mask.dim() == 3:
+        # (B, S_q, S_kv) -> (B, 1, S_q, S_kv)
+        attn_mask = attn_mask.unsqueeze(1)
+    if attn_mask.dim() == 4 and attn_mask.shape[-2] == 1 and kv_len > 1:
+        # (B, 1, 1, S) -> (B, 1, S_q, S)
+        attn_mask = attn_mask.expand(-1, -1, q_len, -1)
+    return attn_mask
+
+
+def _patched_attention_forward(self, x: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
+    b, t, c = x.shape
+
+    q = self.to_heads(self.to_q(x))
+    k = self.to_heads(self.to_k(x))
+    v = self.to_heads(self.to_v(x))
+
+    q = self.q_norm(q)
+    k = self.k_norm(k)
+
+    attn_mask = _expand_attn_mask_for_npu(attn_mask, q_len=t, kv_len=t)
+    x = F.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        attn_mask=attn_mask,
+        dropout_p=self.attn_drop.p if self.training else 0.0,
+    )
+    x = x.transpose(1, 2).reshape(b, t, -1)
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x
+
+
+def _patched_attention_forward_chunk(
+    self,
+    x: torch.Tensor,
+    att_cache: torch.Tensor | None = None,
+    attn_mask: torch.Tensor | None = None,
+):
+    b, t, c = x.shape
+
+    q = self.to_heads(self.to_q(x))
+    k = self.to_heads(self.to_k(x))
+    v = self.to_heads(self.to_v(x))
+
+    q = self.q_norm(q)
+    k = self.k_norm(k)
+
+    if att_cache is not None:
+        k_cache, v_cache = att_cache.chunk(2, dim=3)
+        k = torch.cat([k, k_cache], dim=2)
+        v = torch.cat([v, v_cache], dim=2)
+
+    new_att_cache = torch.cat([k, v], dim=3)
+    kv_len = k.shape[2]
+    if attn_mask is not None:
+        attn_mask = _expand_attn_mask_for_npu(attn_mask, q_len=t, kv_len=kv_len)
+    x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+    x = x.transpose(1, 2).reshape(b, t, -1)
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x, new_att_cache
+
+
+@contextmanager
+def npu_math_sdpa_context() -> Iterator[None]:
+    """Force SDPA MATH backend so Ascend does not call fused FA."""
+    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
+
+
+def _disable_upsample_encoder_compile() -> None:
+    """Run CosyVoice2 UpsampleConformerEncoderV2 eagerly on Ascend.
+
+    ``forward_chunk`` is decorated with ``@torch.compile(dynamic=True,
+    backend="eager")`` upstream. The backend gives no speedup (still eager),
+    but Dynamo tracing of the relative-position attention with a KV cache
+    creates a symbolic add ``matrix_ac(..., chunk+cache) + matrix_bd(...,
+    (pos//2)+1)`` whose sizes Ascend torch's fake-tensor ``infer_size``
+    cannot reconcile, crashing on the 2nd+ streaming chunk. Unwrapping the
+    compile and disabling Dynamo keeps concrete shapes (always equal).
+    """
+    try:
+        from cosyvoice2.transformer import upsample_encoder_v2
+    except ImportError:
+        return
+
+    enc_cls = getattr(upsample_encoder_v2, "UpsampleConformerEncoderV2", None)
+    if enc_cls is None:
+        return
+
+    fn = enc_cls.forward_chunk
+    orig = getattr(fn, "_torchdynamo_orig_callable", None) or getattr(fn, "__wrapped__", fn)
+    enc_cls.forward_chunk = torch._dynamo.disable(orig)  # type: ignore[method-assign]
+    logger.info("Disabled torch.compile on CosyVoice2 UpsampleConformerEncoderV2.forward_chunk (Ascend eager)")
+
+
+def apply_cosyvoice2_dit_attn_npu_patch() -> None:
+    """Monkey-patch CosyVoice2 DiT Attention for Ascend FA mask constraints."""
+    global _PATCHED
+    if _PATCHED:
+        return
+
+    try:
+        from cosyvoice2.flow import decoder_dit
+    except ImportError:
+        logger.debug("cosyvoice2 not installed; skip DiT attn NPU patch")
+        return
+
+    attn_cls = getattr(decoder_dit, "Attention", None)
+    if attn_cls is None:
+        return
+
+    attn_cls.forward = _patched_attention_forward  # type: ignore[method-assign]
+    attn_cls.forward_chunk = _patched_attention_forward_chunk  # type: ignore[method-assign]
+    _disable_upsample_encoder_compile()
+    _PATCHED = True
+    logger.info("Applied CosyVoice2 DiT Attention NPU patch (expand attn_mask to Bx1xSxS for aclnnFlashAttentionScore)")

vllm_omni/platforms/npu/models/step_audio2_token2wav.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NPU patches for Step-Audio2 / MiniCPM Token2Wav.
+
+Ascend-specific workarounds that must not live in the shared GPU model file:
+
+1. HiFT vocoder CPU offload — STFT/ISTFT and sine-source generation are
+   unstable on Ascend (aicore 507015). HiFT is tiny, so running it on CPU
+   is acceptable; inputs/outputs are moved transparently.
+2. CosyVoice2 DiT SDPA — force MATH backend (+ DiT attn mask expand) to
+   avoid fused FA rejecting CosyVoice ``(B,1,1,S)`` masks (error 161001).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager, nullcontext
+from types import MethodType
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_PATCHED = False
+_original_ensure_models_loaded = None
+_original_forward = None
+_original_stream_chunk_for = None
+
+
+def patch_hift_module_for_npu(hift: torch.nn.Module) -> None:
+    """Run the entire HiFT vocoder on CPU when on Ascend NPU.
+
+    HiFT is tiny, so keeping it on NPU yields negligible speedup while
+    several of its operators are unstable on Ascend and trigger aicore
+    exceptions (runtime error 507015): the STFT/ISTFT as well as the
+    sine-source generation (``f0_upsamp`` + ``m_source`` cumsum/sin over
+    the full waveform length). Running the whole module on CPU sidesteps
+    all of them. Inputs are copied to CPU and outputs are moved back to
+    the original device transparently, so callers stay unchanged.
+    """
+    if getattr(hift, "_npu_cpu_offload_patched", False):
+        return
+
+    hift.to("cpu")
+    original_forward = hift.forward
+
+    def _forward_on_cpu(_module, *args, **kwargs):
+        output_device: torch.device | None = None
+
+        def to_cpu(value):
+            nonlocal output_device
+            if isinstance(value, torch.Tensor):
+                if output_device is None and value.device.type != "cpu":
+                    output_device = value.device
+                return value.cpu()
+            return value
+
+        cpu_args = tuple(to_cpu(a) for a in args)
+        cpu_kwargs = {k: to_cpu(v) for k, v in kwargs.items()}
+        output = original_forward(*cpu_args, **cpu_kwargs)
+
+        if output_device is None:
+            return output
+
+        def to_device(value):
+            if isinstance(value, torch.Tensor):
+                return value.to(output_device)
+            return value
+
+        if isinstance(output, tuple):
+            return tuple(to_device(v) for v in output)
+        return to_device(output)
+
+    hift.forward = MethodType(_forward_on_cpu, hift)
+    hift._npu_cpu_offload_patched = True
+
+
+@contextmanager
+def npu_token2wav_sdpa_context() -> 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,
+        )
+
+        apply_cosyvoice2_dit_attn_npu_patch()
+        with npu_math_sdpa_context():
+            yield
+    except Exception:
+        with nullcontext():
+            yield
+
+
+def _patched_ensure_models_loaded(self) -> None:
+    assert _original_ensure_models_loaded is not None
+    was_loaded = self._models_loaded
+    _original_ensure_models_loaded(self)
+    if was_loaded or self.device.type != "npu" or self._hift is None:
+        return
+    patch_hift_module_for_npu(self._hift)
+
+
+def _patched_forward(self, generated_speech_tokens, prompt_wav, return_bytes=True):
+    assert _original_forward is not None
+    if self.device.type != "npu":
+        return _original_forward(self, generated_speech_tokens, prompt_wav, return_bytes)
+    with npu_token2wav_sdpa_context():
+        return _original_forward(self, generated_speech_tokens, prompt_wav, return_bytes)
+
+
+def _patched_stream_chunk_for(self, audio_tokens, prompt_wav, last_chunk, state):
+    assert _original_stream_chunk_for is not None
+    if self.device.type != "npu":
+        return _original_stream_chunk_for(self, audio_tokens, prompt_wav, last_chunk, state)
+    with npu_token2wav_sdpa_context():
+        return _original_stream_chunk_for(self, audio_tokens, prompt_wav, last_chunk, state)
+
+
+def apply_step_audio2_token2wav_npu_patch() -> None:
+    """Monkey-patch StepAudio2Token2WavCore for Ascend NPU.
+
+    Import is deferred and optional: platform bootstrap (e.g. resolving
+    ``current_omni_platform`` from rotary embedding) must not require
+    Token2Wav optional deps such as ``librosa``.
+    """
+    global _PATCHED, _original_ensure_models_loaded, _original_forward, _original_stream_chunk_for
+    if _PATCHED:
+        return
+
+    try:
+        from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import (
+            StepAudio2Token2WavCore,
+        )
+    except ImportError as e:
+        logger.debug("step_audio2 token2wav deps unavailable; skip NPU patch: %s", e)
+        return
+
+    _original_ensure_models_loaded = StepAudio2Token2WavCore._ensure_models_loaded
+    _original_forward = StepAudio2Token2WavCore.forward
+    _original_stream_chunk_for = StepAudio2Token2WavCore.stream_chunk_for
+
+    StepAudio2Token2WavCore._ensure_models_loaded = _patched_ensure_models_loaded  # type: ignore[method-assign]
+    StepAudio2Token2WavCore.forward = _patched_forward  # type: ignore[method-assign]
+    StepAudio2Token2WavCore.stream_chunk_for = _patched_stream_chunk_for  # type: ignore[method-assign]
+
+    _PATCHED = True
+    logger.debug("Applied NPU patch for StepAudio2Token2WavCore")