diff --git a/examples/online_serving/minicpmo/README.md b/examples/online_serving/minicpmo/README.md index a34764652..334b5ed11 100644 --- a/examples/online_serving/minicpmo/README.md +++ b/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 \ diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md index 8b3d2909d..4bbf416ba 100644 --- a/recipes/OpenBMB/MiniCPM-o-4_5.md +++ b/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) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 980615479..d191dddcc 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/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 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml new file mode 100644 index 000000000..7a9c400cc --- /dev/null +++ b/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 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index f5c0b03a6..69f77eb6d 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/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: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index b930578d6..3019a5b3d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/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 ============== diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 083f54afe..d14fe14cf 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/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 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py new file mode 100644 index 000000000..c069b8d09 --- /dev/null +++ b/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(" 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)") diff --git a/vllm_omni/platforms/npu/models/step_audio2_token2wav.py b/vllm_omni/platforms/npu/models/step_audio2_token2wav.py new file mode 100644 index 000000000..04768d04d --- /dev/null +++ b/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")