anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[NPU] Support VoxCPM2 model (#4310)

Browse Source 
Signed-off-by: tanhaoan333 <tanhaoan@huawei.com>
This commit is contained in:
tanhaoan333 authored and GitHub committed 2026-06-11 14:12:06 +08:00
1 parent 0d3c259a76
commit e26f5cb918
2 files changed
+106 -31
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
+22 -16
View File
@@ -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]
vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+84 -15
View File
@@ -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,