diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index a3767df1d..19df5bf89 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -36,6 +36,7 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.utils.speaker_cache import ( get_speaker_cache, iter_custom_voice_profiles, @@ -863,11 +864,11 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): # the checkpoint download/read cost at construction time; DummyModelLoader # will then randomize the just-loaded _tts params — this is intended. model_path = vllm_config.model_config.model + self._device = current_omni_platform.get_torch_device() VoxCPM = import_voxcpm2_core() native = VoxCPM.from_pretrained(model_path, load_denoiser=False, optimize=False) - self._tts: nn.Module = native.tts_model.to("cuda") + self._tts: nn.Module = native.tts_model.to(self._device) self._side_dtype = self._tts.fusion_concat_proj.weight.dtype - self._device = "cuda" self._patch_size = self._tts.patch_size self._feat_dim = self._tts.feat_dim self._sample_rate = getattr(self.config, "sample_rate", 48000) @@ -888,8 +889,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): # for sliding-window streaming decode. 12 matches the nanovllm reference # implementation and covers the longest VAE decoder receptive field. self._n_decode_pad_frames = 12 - self._enable_torch_compile = True - self._compile_vae = True + use_cuda_graph = current_omni_platform.is_cuda() + self._enable_torch_compile = current_omni_platform.supports_torch_inductor() + self._compile_vae = self._enable_torch_compile self._max_decode_steps = 2000 self._max_batch_size = getattr(vllm_config.scheduler_config, "max_num_seqs", 4) @@ -900,15 +902,15 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._enable_profiling = self._runtime_config.enable_profiling self._perf = _PerfTimer(enabled=self._enable_profiling) self._cfm_buffers: _CFMBufferManager | None = None - self._enable_cuda_graph = True + self._enable_cuda_graph = use_cuda_graph self._scaffold_graphs: dict[int, _CapturedGraph] = {} self._residual_graphs: dict[int, _CapturedGraph] = {} self._decode_graph_capture_policy = self._runtime_config.decode_graph_capture_policy self._vae_graphs: dict[tuple[int, int], _CapturedVAEGraph] = {} - self._enable_vae_cuda_graph = self._runtime_config.enable_vae_cuda_graph + self._enable_vae_cuda_graph = use_cuda_graph and self._runtime_config.enable_vae_cuda_graph self._cfm_graphs: dict[int, _CapturedCFMGraph] = {} - self._enable_cfm_cuda_graph = self._runtime_config.enable_cfm_cuda_graph - self._enable_cfm_prealloc_output = self._runtime_config.enable_cfm_prealloc_output + self._enable_cfm_cuda_graph = use_cuda_graph and self._runtime_config.enable_cfm_cuda_graph + self._enable_cfm_prealloc_output = use_cuda_graph and self._runtime_config.enable_cfm_prealloc_output self._enable_batched_cfm = self._runtime_config.enable_batched_cfm self._deterministic_cfm_noise = self._runtime_config.deterministic_cfm_noise self._deterministic_cfm_seed = self._runtime_config.deterministic_cfm_seed @@ -916,7 +918,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._audio_emit_every = self._runtime_config.audio_emit_every self._vae_decode_every = self._runtime_config.vae_decode_every self._enable_delayed_audio_copy = self._runtime_config.enable_delayed_audio_copy - self._delayed_audio_copy_use_events = self._runtime_config.delayed_audio_copy_use_events + self._delayed_audio_copy_use_events = use_cuda_graph and self._runtime_config.delayed_audio_copy_use_events self._coalesce_audio_d2h = self._runtime_config.coalesce_audio_d2h self._enable_batched_vae_decode = self._runtime_config.enable_batched_vae_decode self._enable_batched_fsq_fusion = self._runtime_config.enable_batched_fsq_fusion @@ -1382,7 +1384,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): return sr_cond def _run_vae_decode(self, feat: torch.Tensor) -> torch.Tensor: - if not feat.is_cuda: + if feat.device.type != current_omni_platform.device_type: return self.tts.audio_vae.decode(feat) sr_cond = self._get_vae_decode_sr_cond(feat.device) @@ -1817,7 +1819,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): def _run_cfm(self, dit_h: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: with _NvtxRange("voxcpm2.cfm"): if self._cfm_buffers is not None: - if self._enable_cfm_cuda_graph and dit_h.is_cuda: + if self._enable_cfm_cuda_graph and dit_h.device.type == current_omni_platform.device_type: return self._run_cfm_cuda_graph(dit_h, cond).transpose(1, 2) return _optimized_solve_euler( self.tts.feat_decoder, @@ -1841,7 +1843,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): def _run_cfm_for_state(self, state: _RequestState, dit_h: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: with _NvtxRange("voxcpm2.cfm"): if self._cfm_buffers is not None: - if self._enable_cfm_cuda_graph and dit_h.is_cuda: + if self._enable_cfm_cuda_graph and dit_h.device.type == current_omni_platform.device_type: if self._deterministic_cfm_noise and not self._enable_cfm_prealloc_output: graph = self._get_cfm_cuda_graph(dit_h, cond) with _NvtxRange("voxcpm2.cfm.graph_copy_mu"): @@ -2057,7 +2059,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): def _enqueue_delayed_audio_copy(self, state: _RequestState, audio: torch.Tensor) -> None: src = audio.detach().contiguous() - if not src.is_cuda: + if src.device.type != current_omni_platform.device_type: state.pending_audio_copies.append(_PendingAudioCopy(host=src.cpu().contiguous())) return @@ -2145,7 +2147,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if state.is_stopping or state.precomputed_is_stopping is not None: continue stop_logits = state.precomputed_stop_logits - if stop_logits is None or not stop_logits.is_cuda: + if stop_logits is None or stop_logits.device.type != current_omni_platform.device_type: continue pending.append((state, stop_logits)) if not pending: @@ -2426,7 +2428,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if self._uses_sparse_audio_outputs(): ready_req_ids = list(audio_by_req) chunks = [audio_by_req[req_id].reshape(-1) for req_id in ready_req_ids] - if self._coalesce_audio_d2h and any(chunk.is_cuda for chunk in chunks): + if self._coalesce_audio_d2h and any( + chunk.device.type == current_omni_platform.device_type for chunk in chunks + ): sizes = [int(chunk.numel()) for chunk in chunks] merged = torch.cat(chunks, dim=0) if len(chunks) > 1 else chunks[0] merged_cpu = merged.detach().cpu().contiguous() @@ -2435,7 +2439,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): mm["model_outputs"] = chunks mm["sr"] = [sr for _ in ready_req_ids] mm["meta"] = {"req_id": ready_req_ids, "sparse_audio": ["1"]} - elif self._coalesce_audio_d2h and any(audio.is_cuda for audio in audio_by_req.values()): + elif self._coalesce_audio_d2h and any( + audio.device.type == current_omni_platform.device_type for audio in audio_by_req.values() + ): ready_req_ids = list(audio_by_req) chunks = [audio_by_req[req_id].reshape(-1) for req_id in ready_req_ids] sizes = [int(chunk.numel()) for chunk in chunks] diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index ed49ddee5..9ea236bb7 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -240,9 +240,38 @@ class NPUARModelRunner(OmniNPUModelRunner): def _resolve_pooler_payload_req_ids(self, req_ids_output_copy: list[str]) -> tuple[str, list[str]]: downstream_req_ids = [rid for rid in req_ids_output_copy if self._request_needs_downstream_stage_payload(rid)] engine_output_type = (self.vllm_config.model_config.engine_output_type or "").lower() + # Single-stage AR TTS models (e.g. VoxCPM2) finish on this stage but still + # need multimodal payloads for final audio postprocess/output. if engine_output_type == "audio" and not downstream_req_ids: downstream_req_ids = req_ids_output_copy return engine_output_type, downstream_req_ids + + @staticmethod + def _sparse_mm_req_ids(multimodal_outputs: Any) -> list[str] | None: + if not isinstance(multimodal_outputs, dict): + return None + meta = multimodal_outputs.get("meta") + req_ids = None + sparse_audio = False + if isinstance(meta, dict): + req_ids = meta.get("req_id") + sparse_audio = NPUARModelRunner._is_sparse_audio_marker(meta.get("sparse_audio")) + if req_ids is None: + req_ids = multimodal_outputs.get("meta.req_id") + sparse_audio = NPUARModelRunner._is_sparse_audio_marker(multimodal_outputs.get("meta.sparse_audio")) + if not sparse_audio: + return None + if not isinstance(req_ids, list): + return None + return [rid for rid in req_ids if isinstance(rid, str)] + + @staticmethod + def _is_sparse_audio_marker(value: Any) -> bool: + if isinstance(value, list): + return any(str(item).lower() in ("1", "true", "yes", "on") for item in value) + if isinstance(value, str): + return value.lower() in ("1", "true", "yes", "on") + return bool(value) # -------------------------------------- Omni-new ------------------------------------------------- @torch.inference_mode() @@ -910,16 +939,27 @@ class NPUARModelRunner(OmniNPUModelRunner): # -------------------------------------- Omni-new ------------------------------------------------- engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) + sparse_mm_req_ids = self._sparse_mm_req_ids(multimodal_outputs) + sparse_mm_index = {rid: i for i, rid in enumerate(sparse_mm_req_ids or [])} + if engine_output_type == "audio" and sparse_mm_req_ids is not None: + sparse_req_id_set = set(sparse_mm_req_ids) + downstream_req_ids = [rid for rid in req_ids_output_copy if rid in sparse_req_id_set] needs_pooler_payload = len(downstream_req_ids) > 0 downstream_req_id_set = set(downstream_req_ids) hidden_states_cpu = None req_hidden_states_cpu: dict[str, torch.Tensor] | None = None - if needs_pooler_payload: + audio_sparse_output = engine_output_type == "audio" and sparse_mm_req_ids is not None + needs_scheduled_hidden_payload = needs_pooler_payload and ( + self.omni_prefix_cache is None or not self._model_needs_full_prefix_hidden_states() + ) + if needs_scheduled_hidden_payload: num_valid_tokens = min( int(scheduler_output.total_num_scheduled_tokens), int(hidden_states.shape[0]), ) - if len(downstream_req_ids) == len(req_ids_output_copy): + if audio_sparse_output: + pass + elif len(downstream_req_ids) == len(req_ids_output_copy): hidden_states_cpu = hidden_states[:num_valid_tokens].detach().to("cpu").contiguous() else: req_hidden_states_cpu = {} @@ -931,9 +971,14 @@ class NPUARModelRunner(OmniNPUModelRunner): dtype=np.int32, ) query_start_loc_cpu = self.query_start_loc.cpu + if callable(query_start_loc_cpu): + query_start_loc_cpu = query_start_loc_cpu() pooler_output: list[dict[str, object]] | None = None if needs_pooler_payload: + combined_hidden_states = None + combined_multimodal_outputs = None + mm_cpu = None if self.omni_prefix_cache is not None: ( combined_hidden_states, @@ -943,8 +988,10 @@ class NPUARModelRunner(OmniNPUModelRunner): multimodal_outputs, scheduler_output.num_scheduled_tokens, ) - else: - mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs)) + if self.omni_prefix_cache is None or combined_multimodal_outputs is None: + mm_cpu = build_mm_cpu( + flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs + ) self._process_additional_information_updates( hidden_states, @@ -973,17 +1020,20 @@ class NPUARModelRunner(OmniNPUModelRunner): start = int(query_start_loc_cpu[idx]) sched = int(num_scheduled_tokens_np[idx]) end = start + sched - if req_hidden_states_cpu is not None and combined_hidden_states is None: - req_hidden_states = req_hidden_states_cpu[rid] - else: - req_hidden_states = self._resolve_req_hidden_states( - hidden_states_cpu, - combined_hidden_states, - rid, - start, - end, - ) - payload: dict[str, object] = {"hidden": req_hidden_states} + payload: dict[str, object] = {} + if not audio_sparse_output: + if req_hidden_states_cpu is not None and combined_hidden_states is None: + req_hidden_states = req_hidden_states_cpu[rid] + else: + req_hidden_states = self._resolve_req_hidden_states( + hidden_states_cpu, + combined_hidden_states, + rid, + start, + end, + ) + payload["hidden"] = req_hidden_states + mm_payload: dict[str, object] = {} if combined_multimodal_outputs or mm_cpu: if combined_multimodal_outputs: @@ -1006,6 +1056,25 @@ class NPUARModelRunner(OmniNPUModelRunner): mm_payload[mm_key] = _unwrap_lists(combined_multimodal_outputs[mm_key][rid]) else: for mm_key, mm_val in mm_cpu.items(): + if mm_key in {"meta.req_id", "meta.sparse_audio"}: + continue + if audio_sparse_output and isinstance(mm_val, list): + sparse_idx = sparse_mm_index.get(rid) + if sparse_idx is None: + continue + if sparse_idx >= len(mm_val): + logger.warning( + "Sparse multimodal payload mismatch for request %s: index %d >= %d.", + rid, + sparse_idx, + len(mm_val), + ) + continue + sparse_val = mm_val[sparse_idx] + mm_payload[mm_key] = ( + sparse_val.clone() if isinstance(sparse_val, torch.Tensor) else sparse_val + ) + continue mm_payload[mm_key] = to_payload_element( element=mm_val, idx=idx,