1327 lines
55 KiB
Python
1327 lines
55 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""VoxCPM2 AR talker — PagedAttention pipeline with per-request state.
|
|
|
|
Architecture:
|
|
MiniCPM4PagedForVoxCPM2 (base_lm, 28 layers, PagedAttention + fp32 RoPE)
|
|
→ FSQ → MiniCPM4PagedResidualLM (8 layers, PagedAttention, no RoPE)
|
|
→ LocDiT (CFM solver) → AudioVAE → 48kHz waveform
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import dataclasses
|
|
import logging
|
|
import math
|
|
import os
|
|
import time
|
|
from collections.abc import Iterable
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from vllm.config import VllmConfig
|
|
from vllm.forward_context import get_forward_context, override_forward_context
|
|
from vllm.inputs import tokens_input
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.models.utils import (
|
|
AutoWeightsLoader,
|
|
WeightsMapper,
|
|
maybe_prefix,
|
|
)
|
|
from vllm.multimodal.audio import AudioResampler
|
|
from vllm.sequence import IntermediateTensors
|
|
|
|
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
|
from vllm_omni.utils.speaker_cache import get_speaker_cache
|
|
|
|
from .minicpm4_paged import MiniCPM4PagedForVoxCPM2, MiniCPM4PagedResidualLM
|
|
from .voxcpm2_import_utils import import_voxcpm2_core
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_ENABLE_PROFILING = os.environ.get("VOXCPM2_PROFILE", "0") == "1"
|
|
|
|
# Lower bound for the _active_states leak-warn threshold. The effective
|
|
# threshold is max(_ACTIVE_STATE_LEAK_WARN_MIN, 4 * max_batch_size) so small
|
|
# deployments still get a usable floor instead of a tiny noisy one.
|
|
_ACTIVE_STATE_LEAK_WARN_MIN = 512
|
|
|
|
|
|
def is_cjk_char(c: str) -> bool:
|
|
"""Check if a character is a CJK ideograph."""
|
|
cp = ord(c)
|
|
return (
|
|
0x4E00 <= cp <= 0x9FFF # CJK Unified Ideographs
|
|
or 0x3400 <= cp <= 0x4DBF # Extension A
|
|
or 0xF900 <= cp <= 0xFAFF # Compatibility Ideographs
|
|
or 0x20000 <= cp <= 0x2A6DF # Extension B
|
|
or 0x2A700 <= cp <= 0x2B73F # Extension C
|
|
or 0x2B740 <= cp <= 0x2B81F # Extension D
|
|
or 0x2F800 <= cp <= 0x2FA1F # Compatibility Supplement
|
|
)
|
|
|
|
|
|
def build_cjk_split_map(tokenizer: Any) -> dict[int, list[int]]:
|
|
"""Build {multichar_cjk_token_id: [single_char_ids]} from tokenizer vocab."""
|
|
vocab = tokenizer.get_vocab()
|
|
split_map: dict[int, list[int]] = {}
|
|
for token, token_id in vocab.items():
|
|
clean = token.replace("\u2581", "")
|
|
if len(clean) >= 2 and all(is_cjk_char(c) for c in clean):
|
|
char_ids = tokenizer.convert_tokens_to_ids(list(clean))
|
|
if all(cid != tokenizer.unk_token_id for cid in char_ids):
|
|
split_map[token_id] = char_ids
|
|
return split_map
|
|
|
|
|
|
def split_multichar_chinese(token_ids: list[int], split_map: dict[int, list[int]]) -> list[int]:
|
|
"""Replace multichar Chinese token IDs with single-char IDs (idempotent)."""
|
|
result: list[int] = []
|
|
for tid in token_ids:
|
|
expansion = split_map.get(tid)
|
|
if expansion is not None:
|
|
result.extend(expansion)
|
|
else:
|
|
result.append(tid)
|
|
return result
|
|
|
|
|
|
def build_voxcpm2_prompt(
|
|
hf_config: Any,
|
|
tokenizer: Any,
|
|
split_map: dict[int, list[int]],
|
|
text: str,
|
|
ref_audio: Any | None = None,
|
|
ref_sr: int | None = None,
|
|
ref_text: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Build a VoxCPM2 prefill prompt whose ``prompt_token_ids`` length matches
|
|
the talker-side prefill length.
|
|
|
|
Used by both online serving (``serving_speech._build_voxcpm2_prompt``) and
|
|
the offline example, so the talker-side length assertion never fires.
|
|
"""
|
|
ids = split_multichar_chinese(tokenizer.encode(text, add_special_tokens=True), split_map)
|
|
bos = tokenizer.bos_token_id
|
|
if ids and ids[0] == bos:
|
|
ids = ids[1:]
|
|
prefill_len = len(ids) + 1 # + audio_start
|
|
additional: dict[str, Any] = {"text_token_ids": [ids]}
|
|
if ref_audio is not None:
|
|
vae = hf_config.audio_vae_config
|
|
patch_samples = hf_config.patch_size * math.prod(vae["encoder_rates"])
|
|
ref_len = math.ceil(math.ceil(len(ref_audio) * vae["sample_rate"] / ref_sr) / patch_samples)
|
|
if ref_text is not None:
|
|
additional["prompt_audio"] = [[ref_audio, ref_sr]]
|
|
additional["prompt_text"] = [ref_text]
|
|
ref_ids = split_multichar_chinese(tokenizer.encode(ref_text, add_special_tokens=True), split_map)
|
|
if ref_ids and ref_ids[0] == bos:
|
|
ref_ids = ref_ids[1:]
|
|
prefill_len += ref_len + len(ref_ids)
|
|
else:
|
|
additional["reference_audio"] = [[ref_audio, ref_sr]]
|
|
prefill_len += ref_len + 2 # ref_start / ref_end
|
|
prompt = tokens_input(prompt_token_ids=[1] * prefill_len)
|
|
prompt["additional_information"] = additional
|
|
return prompt
|
|
|
|
|
|
def _encode_raw_audio(
|
|
tts: nn.Module,
|
|
samples: list[float] | torch.Tensor,
|
|
sr: int,
|
|
padding_mode: str = "right",
|
|
) -> torch.Tensor:
|
|
"""Encode raw audio samples using the native VoxCPM2 AudioVAE.
|
|
|
|
Mirrors ``VoxCPM2Model._encode_wav`` but accepts in-memory samples
|
|
instead of a file path (needed for the OpenAI speech API).
|
|
"""
|
|
if isinstance(samples, list):
|
|
audio = torch.tensor(samples, dtype=torch.float32)
|
|
else:
|
|
audio = samples.float()
|
|
if audio.ndim == 1:
|
|
audio = audio.unsqueeze(0)
|
|
|
|
encode_sr = tts._encode_sample_rate
|
|
if sr != encode_sr:
|
|
audio_np = audio.squeeze(0).numpy()
|
|
resampler = AudioResampler(target_sr=encode_sr)
|
|
audio_np = resampler.resample(audio_np, orig_sr=sr)
|
|
audio = torch.from_numpy(audio_np).unsqueeze(0)
|
|
|
|
patch_len = tts.patch_size * tts.chunk_size
|
|
if audio.size(1) % patch_len != 0:
|
|
padding_size = patch_len - audio.size(1) % patch_len
|
|
pad = (padding_size, 0) if padding_mode == "left" else (0, padding_size)
|
|
audio = torch.nn.functional.pad(audio, pad)
|
|
|
|
feat = tts.audio_vae.encode(audio.to(tts.device), encode_sr).cpu()
|
|
return feat.view(tts.audio_vae.latent_dim, -1, tts.patch_size).permute(1, 2, 0)
|
|
|
|
|
|
# ===================================================================
|
|
# Per-request state
|
|
# ===================================================================
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class _RequestState:
|
|
request_id: str
|
|
curr_embed_for_next: torch.Tensor | None = None
|
|
prev_feat_embed: torch.Tensor | None = None
|
|
curr_prefix_feat_cond: torch.Tensor | None = None
|
|
last_audio_patch_gpu: torch.Tensor | None = None
|
|
precomputed_stop_logits: torch.Tensor | None = None
|
|
# Rolling tail of previously-decoded latents used as VAE receptive-field context.
|
|
# Shape (n_pad_frames, feat_dim) on GPU. None before first decode.
|
|
decode_pad: torch.Tensor | None = None
|
|
decode_step_count: int = 0
|
|
request_start_time: float = 0.0
|
|
prefill_completed: bool = False
|
|
prefill_text: str = ""
|
|
prompt_cache: dict | None = None
|
|
prefill_masks: tuple | None = None
|
|
is_stopping: bool = False
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class _CapturedGraph:
|
|
graph: torch.cuda.CUDAGraph
|
|
input_embeds: torch.Tensor
|
|
positions: torch.Tensor
|
|
output: torch.Tensor
|
|
|
|
|
|
# ===================================================================
|
|
# Profiling timer
|
|
# ===================================================================
|
|
|
|
|
|
class _PerfTimer:
|
|
__slots__ = ("_enabled", "_timers", "_counts", "_starts", "_pairs")
|
|
|
|
def __init__(self, enabled: bool = False):
|
|
self._enabled = enabled
|
|
self._timers: dict[str, float] = {}
|
|
self._counts: dict[str, int] = {}
|
|
self._starts: dict[str, torch.cuda.Event] = {}
|
|
self._pairs: list[tuple[str, torch.cuda.Event, torch.cuda.Event]] = []
|
|
|
|
def start(self, name: str) -> None:
|
|
if not self._enabled:
|
|
return
|
|
evt = torch.cuda.Event(enable_timing=True)
|
|
evt.record()
|
|
self._starts[name] = evt
|
|
|
|
def stop(self, name: str) -> None:
|
|
if not self._enabled or name not in self._starts:
|
|
return
|
|
start_evt = self._starts.pop(name)
|
|
end_evt = torch.cuda.Event(enable_timing=True)
|
|
end_evt.record()
|
|
self._pairs.append((name, start_evt, end_evt))
|
|
|
|
def _resolve(self) -> None:
|
|
if not self._pairs:
|
|
return
|
|
torch.accelerator.synchronize()
|
|
for name, s, e in self._pairs:
|
|
self._timers[name] = self._timers.get(name, 0.0) + s.elapsed_time(e)
|
|
self._counts[name] = self._counts.get(name, 0) + 1
|
|
self._pairs.clear()
|
|
|
|
def breakdown(self) -> str:
|
|
if not self._enabled:
|
|
return ""
|
|
self._resolve()
|
|
if not self._timers:
|
|
return ""
|
|
total = self._timers.get("decode_step", sum(self._timers.values()))
|
|
lines = [
|
|
"=== VoxCPM2 Decode Step Breakdown ===",
|
|
f"{'Component':<30} | {'ms':>10} | {'%':>6} | {'N':>5} | {'avg':>8}",
|
|
"-" * 70,
|
|
]
|
|
for name in sorted(self._timers):
|
|
t, c = self._timers[name], self._counts[name]
|
|
lines.append(f"{name:<30} | {t:>10.2f} | {t / total * 100:>5.1f}% | {c:>5} | {t / c:>8.3f}")
|
|
lines.append(f"{'TOTAL':<30} | {total:>10.2f} |")
|
|
return "\n".join(lines)
|
|
|
|
def reset(self) -> None:
|
|
self._timers.clear()
|
|
self._counts.clear()
|
|
self._starts.clear()
|
|
self._pairs.clear()
|
|
|
|
|
|
# ===================================================================
|
|
# CFM pre-allocated buffers + optimized Euler solver
|
|
# ===================================================================
|
|
|
|
|
|
class _CFMBufferManager:
|
|
def __init__(
|
|
self,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
feat_dim: int,
|
|
patch_size: int,
|
|
dit_hidden_size: int,
|
|
max_batch_size: int = 1,
|
|
sway_sampling_coef: float = 1.0,
|
|
):
|
|
n = 2 * max_batch_size # CFG doubles the batch
|
|
self.x_in = torch.zeros(n, feat_dim, patch_size, device=device, dtype=dtype)
|
|
self.mu_in = torch.zeros(n, dit_hidden_size, device=device, dtype=dtype)
|
|
self.t_in = torch.zeros(n, device=device, dtype=dtype)
|
|
self.dt_in = torch.zeros(n, device=device, dtype=dtype)
|
|
self.cond_in = torch.zeros(n, feat_dim, patch_size, device=device, dtype=dtype)
|
|
self.noise = torch.zeros(max_batch_size, feat_dim, patch_size, device=device, dtype=dtype)
|
|
self._sway_coef = sway_sampling_coef
|
|
self._device = device
|
|
self._dtype = dtype
|
|
self.t_span_10 = self._make_t_span(10)
|
|
|
|
def _make_t_span(self, n: int) -> torch.Tensor:
|
|
t = torch.linspace(1, 0, n + 1, device=self._device, dtype=self._dtype)
|
|
return t + self._sway_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
|
|
|
|
def get_t_span(self, n: int) -> torch.Tensor:
|
|
return self.t_span_10 if n == 10 else self._make_t_span(n)
|
|
|
|
|
|
def _optimized_solve_euler(
|
|
cfm_module: nn.Module,
|
|
mu: torch.Tensor,
|
|
patch_size: int,
|
|
cond: torch.Tensor,
|
|
n_timesteps: int,
|
|
cfg_value: float,
|
|
buffers: _CFMBufferManager,
|
|
use_cfg_zero_star: bool = True,
|
|
cfg_cutoff_ratio: float = 1.0,
|
|
perf: _PerfTimer | None = None,
|
|
) -> torch.Tensor:
|
|
estimator = cfm_module.estimator
|
|
mean_mode = getattr(cfm_module, "mean_mode", False)
|
|
b = mu.size(0)
|
|
|
|
buffers.noise[:b].normal_()
|
|
x = buffers.noise[:b].clone()
|
|
|
|
t_span = buffers.get_t_span(n_timesteps)
|
|
t, dt = t_span[0], t_span[0] - t_span[1]
|
|
zero_init_steps = max(1, int(len(t_span) * 0.04))
|
|
cfg_cutoff_step = max(zero_init_steps + 1, int(len(t_span) * cfg_cutoff_ratio))
|
|
|
|
for step in range(1, len(t_span)):
|
|
if use_cfg_zero_star and step <= zero_init_steps:
|
|
dphi_dt = torch.zeros_like(x)
|
|
elif step <= cfg_cutoff_step:
|
|
buffers.x_in[:b].copy_(x)
|
|
buffers.x_in[b : 2 * b].copy_(x)
|
|
buffers.mu_in[:b].copy_(mu)
|
|
buffers.mu_in[b : 2 * b].zero_()
|
|
# Broadcast the 0-dim GPU scalar directly instead of
|
|
# ``.fill_(t.item())`` — ``.item()`` forces a GPU->CPU sync.
|
|
buffers.t_in[: 2 * b].copy_(t)
|
|
if mean_mode:
|
|
buffers.dt_in[: 2 * b].copy_(dt)
|
|
else:
|
|
buffers.dt_in.zero_()
|
|
buffers.cond_in[:b].copy_(cond[:b])
|
|
buffers.cond_in[b : 2 * b].copy_(cond[:b])
|
|
|
|
if perf:
|
|
perf.start(" cfm.estimator_cfg")
|
|
raw_out = estimator(
|
|
buffers.x_in[: 2 * b],
|
|
buffers.mu_in[: 2 * b],
|
|
buffers.t_in[: 2 * b],
|
|
buffers.cond_in[: 2 * b],
|
|
buffers.dt_in[: 2 * b],
|
|
)
|
|
if perf:
|
|
perf.stop(" cfm.estimator_cfg")
|
|
|
|
dphi_dt, cfg_dphi_dt = raw_out[:b], raw_out[b : 2 * b]
|
|
if use_cfg_zero_star:
|
|
pos = dphi_dt.reshape(b, -1)
|
|
neg = cfg_dphi_dt.reshape(b, -1)
|
|
st = torch.sum(pos * neg, 1, keepdim=True) / (torch.sum(neg**2, 1, keepdim=True) + 1e-8)
|
|
st = st.view(b, *([1] * (len(dphi_dt.shape) - 1)))
|
|
else:
|
|
st = 1.0
|
|
dphi_dt = cfg_dphi_dt * st + cfg_value * (dphi_dt - cfg_dphi_dt * st)
|
|
else:
|
|
buffers.x_in[:b].copy_(x)
|
|
buffers.mu_in[:b].copy_(mu)
|
|
# Broadcast the 0-dim GPU scalar; ``.fill_(t.item())`` would sync.
|
|
buffers.t_in[:b].copy_(t)
|
|
if mean_mode:
|
|
buffers.dt_in[:b].copy_(dt)
|
|
else:
|
|
buffers.dt_in[:b].zero_()
|
|
buffers.cond_in[:b].copy_(cond[:b])
|
|
if perf:
|
|
perf.start(" cfm.estimator_nocfg")
|
|
dphi_dt = estimator(
|
|
buffers.x_in[:b], buffers.mu_in[:b], buffers.t_in[:b], buffers.cond_in[:b], buffers.dt_in[:b]
|
|
)
|
|
if perf:
|
|
perf.stop(" cfm.estimator_nocfg")
|
|
|
|
x = x - dt * dphi_dt
|
|
t = t - dt
|
|
if step < len(t_span) - 1:
|
|
dt = t - t_span[step + 1]
|
|
return x
|
|
|
|
|
|
# ===================================================================
|
|
# Main talker model
|
|
# ===================================================================
|
|
|
|
|
|
class VoxCPM2TalkerForConditionalGeneration(nn.Module):
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
super().__init__()
|
|
self.vllm_config = vllm_config
|
|
self.config = vllm_config.model_config.hf_config
|
|
|
|
self.have_multimodal_outputs = True
|
|
self.has_preprocess = True
|
|
self.has_postprocess = True
|
|
|
|
self.model = MiniCPM4PagedForVoxCPM2(
|
|
vllm_config=vllm_config,
|
|
prefix=maybe_prefix(prefix, "model"),
|
|
)
|
|
self.residual_model = MiniCPM4PagedResidualLM(
|
|
vllm_config=vllm_config,
|
|
prefix=maybe_prefix(prefix, "residual_model"),
|
|
)
|
|
self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
|
|
|
|
# Eager-init tts_model so it registers in self.state_dict() before vLLM's
|
|
# post-__init__ profiling claims the remaining GPU memory for KV cache.
|
|
# Required for load_format=dummy: DummyModelLoader only randomizes
|
|
# already-registered nn.Parameters.
|
|
# NOTE: from_pretrained() is unconditional, so load_format=dummy still pays
|
|
# 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
|
|
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._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)
|
|
|
|
# base_lm/residual_lm in native tts_model duplicate self.model and
|
|
# self.residual_model: copy residual weights over, drop both submodules.
|
|
self.residual_model.load_weights_from_native(self._tts.residual_lm)
|
|
del self._tts.base_lm
|
|
self._tts.base_lm = None
|
|
del self._tts.residual_lm
|
|
self._tts.residual_lm = None
|
|
torch.accelerator.empty_cache()
|
|
|
|
self._inference_timesteps = 10
|
|
self._cfg_value = 2.0
|
|
self._cfg_cutoff_ratio = 1.0
|
|
# Number of trailing latent frames to keep as VAE receptive-field context
|
|
# 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
|
|
self._max_decode_steps = 2000
|
|
self._max_batch_size = getattr(vllm_config.scheduler_config, "max_num_seqs", 4)
|
|
|
|
# Speaker cache for ref_audio_feat across requests
|
|
self._speaker_cache = get_speaker_cache()
|
|
|
|
self._perf = _PerfTimer(enabled=_ENABLE_PROFILING)
|
|
self._cfm_buffers: _CFMBufferManager | None = None
|
|
self._enable_cuda_graph = True
|
|
self._scaffold_graphs: dict[int, _CapturedGraph] = {}
|
|
self._residual_graphs: dict[int, _CapturedGraph] = {}
|
|
self._max_cached_graphs = self._max_batch_size
|
|
self._cuda_graph_pool: tuple | None = None
|
|
self._cuda_graph_warmup_steps = 0
|
|
self._cuda_graph_warmup_threshold = 3
|
|
|
|
self._multichar_zh_split: dict[int, list[int]] | None = None
|
|
|
|
self._active_states: dict[str, _RequestState] = {}
|
|
self._current_request_id: str | None = None
|
|
self._pending_requests: list[tuple[str, bool, torch.Tensor | None, int]] = []
|
|
self._results_queue: list[tuple[str, torch.Tensor | None]] = []
|
|
self._audio_queue: list[tuple[str, Any]] = []
|
|
self._deferred_cleanup_ids: set[str] = set()
|
|
self._active_state_warn_threshold = max(_ACTIVE_STATE_LEAK_WARN_MIN, 4 * self._max_batch_size)
|
|
# one-shot by design: fires at most once per process to avoid log spam.
|
|
self._active_state_warned = False
|
|
|
|
@property
|
|
def tts(self) -> nn.Module:
|
|
return self._tts
|
|
|
|
# -------------------- request state management --------------------
|
|
|
|
def _get_or_create_state(self, request_id: str) -> _RequestState:
|
|
state = self._active_states.get(request_id)
|
|
if state is None:
|
|
state = _RequestState(request_id=request_id)
|
|
self._active_states[request_id] = state
|
|
if len(self._active_states) > self._active_state_warn_threshold and not self._active_state_warned:
|
|
logger.warning(
|
|
"VoxCPM2: _active_states size=%d exceeds threshold %d "
|
|
"(max_batch_size=%d); possible cleanup path leak",
|
|
len(self._active_states),
|
|
self._active_state_warn_threshold,
|
|
self._max_batch_size,
|
|
)
|
|
self._active_state_warned = True
|
|
return state
|
|
|
|
def _switch_to_request(self, request_id: str) -> _RequestState:
|
|
if request_id != self._current_request_id:
|
|
self._current_request_id = request_id
|
|
return self._get_or_create_state(request_id)
|
|
|
|
def _cleanup_request(self, request_id: str) -> None:
|
|
self._active_states.pop(request_id, None)
|
|
if self._current_request_id == request_id:
|
|
self._current_request_id = None
|
|
|
|
def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
|
|
# Defer cleanup: on_requests_finished is called before forward(),
|
|
# so we must not delete state that the current step may still need.
|
|
self._deferred_cleanup_ids.update(finished_req_ids)
|
|
|
|
def _flush_deferred_cleanup(self) -> None:
|
|
for req_id in self._deferred_cleanup_ids:
|
|
self._cleanup_request(req_id)
|
|
self._deferred_cleanup_ids.clear()
|
|
|
|
def _build_prompt_cache(
|
|
self,
|
|
ref_audio: Any = None,
|
|
prompt_audio: Any = None,
|
|
prompt_text: str | None = None,
|
|
) -> dict | None:
|
|
"""Build prompt cache, handling both file paths and raw audio data.
|
|
|
|
The OpenAI speech API sends decoded audio as [samples_list, sr]
|
|
via ``_resolve_ref_audio``, while offline usage sends file paths.
|
|
"""
|
|
tts = self.tts
|
|
|
|
def _is_raw_audio(v: Any) -> bool:
|
|
import numbers
|
|
|
|
return (
|
|
isinstance(v, (list, tuple))
|
|
and len(v) == 2
|
|
and isinstance(v[1], numbers.Integral)
|
|
and isinstance(v[0], (list, torch.Tensor))
|
|
)
|
|
|
|
if not _is_raw_audio(ref_audio) and not _is_raw_audio(prompt_audio):
|
|
return tts.build_prompt_cache(
|
|
prompt_text=prompt_text,
|
|
prompt_wav_path=prompt_audio,
|
|
reference_wav_path=ref_audio,
|
|
)
|
|
|
|
cache: dict[str, Any] = {}
|
|
if ref_audio is not None:
|
|
if _is_raw_audio(ref_audio):
|
|
samples, sr = ref_audio
|
|
cache["ref_audio_feat"] = _encode_raw_audio(tts, samples, sr)
|
|
else:
|
|
cache["ref_audio_feat"] = tts._encode_wav(ref_audio, padding_mode="right")
|
|
|
|
if prompt_audio is not None and prompt_text is not None:
|
|
cache["prompt_text"] = prompt_text
|
|
if _is_raw_audio(prompt_audio):
|
|
samples, sr = prompt_audio
|
|
cache["audio_feat"] = _encode_raw_audio(tts, samples, sr, padding_mode="left")
|
|
else:
|
|
cache["audio_feat"] = tts._encode_wav(prompt_audio, padding_mode="left")
|
|
|
|
has_ref = "ref_audio_feat" in cache
|
|
has_prompt = "audio_feat" in cache
|
|
if has_ref and has_prompt:
|
|
cache["mode"] = "ref_continuation"
|
|
elif has_ref:
|
|
cache["mode"] = "reference"
|
|
else:
|
|
cache["mode"] = "continuation"
|
|
|
|
return cache
|
|
|
|
# -------------------- compile setup --------------------
|
|
|
|
def _setup_cfm_buffers(self) -> None:
|
|
if self._cfm_buffers is not None:
|
|
return
|
|
tts = self.tts
|
|
dit_hidden = tts.lm_to_dit_proj.out_features + tts.res_to_dit_proj.out_features
|
|
self._cfm_buffers = _CFMBufferManager(
|
|
device=torch.device(self._device),
|
|
dtype=self._side_dtype,
|
|
feat_dim=self._feat_dim,
|
|
patch_size=self._patch_size,
|
|
dit_hidden_size=dit_hidden,
|
|
max_batch_size=self._max_batch_size,
|
|
)
|
|
|
|
def _setup_torch_compile(self) -> None:
|
|
if not self._enable_torch_compile:
|
|
return
|
|
tts = self.tts
|
|
estimator = tts.feat_decoder.estimator
|
|
if hasattr(estimator, "_compiled"):
|
|
return
|
|
|
|
targets: list[str] = []
|
|
|
|
try:
|
|
tts.feat_decoder.estimator = torch.compile(estimator, mode="reduce-overhead", fullgraph=False)
|
|
tts.feat_decoder.estimator._compiled = True
|
|
targets.append("LocDiT")
|
|
except Exception as e:
|
|
logger.warning("torch.compile LocDiT failed: %s", e)
|
|
|
|
try:
|
|
if not hasattr(tts.feat_encoder, "_compiled"):
|
|
tts.feat_encoder = torch.compile(tts.feat_encoder, mode="reduce-overhead", fullgraph=False)
|
|
tts.feat_encoder._compiled = True
|
|
targets.append("feat_encoder")
|
|
except Exception as e:
|
|
logger.warning("torch.compile feat_encoder failed: %s", e)
|
|
|
|
if self._compile_vae:
|
|
try:
|
|
if not hasattr(tts.audio_vae, "_compiled"):
|
|
tts.audio_vae.decode = torch.compile(tts.audio_vae.decode, mode="reduce-overhead", fullgraph=False)
|
|
tts.audio_vae._compiled = True
|
|
targets.append("AudioVAE")
|
|
except Exception as e:
|
|
logger.warning("torch.compile AudioVAE failed: %s", e)
|
|
|
|
if not self._enable_cuda_graph:
|
|
if not getattr(self.model, "_selective_compiled", False):
|
|
try:
|
|
targets.extend(f"scaffold.{t}" for t in self.model.compile_selective())
|
|
self.model._selective_compiled = True
|
|
except Exception as e:
|
|
logger.warning("scaffold compile failed: %s", e)
|
|
|
|
if not getattr(self.residual_model, "_selective_compiled", False):
|
|
try:
|
|
targets.extend(f"residual.{t}" for t in self.residual_model.compile_selective())
|
|
self.residual_model._selective_compiled = True
|
|
except Exception as e:
|
|
logger.warning("residual compile failed: %s", e)
|
|
else:
|
|
self.model.precompute_fused_qkv()
|
|
self.residual_model.precompute_fused_qkv()
|
|
targets.append("scaffold+residual (CUDA Graph, skipping compile)")
|
|
|
|
if not getattr(self, "_projections_compiled", False):
|
|
try:
|
|
self._compiled_dit_proj = torch.compile(self._dit_proj_fn, mode="default", fullgraph=True)
|
|
self._compiled_stop_fn = torch.compile(self._stop_fn, mode="default", fullgraph=True)
|
|
self._projections_compiled = True
|
|
targets.append("projections")
|
|
except Exception as e:
|
|
self._compiled_dit_proj = self._compiled_stop_fn = None
|
|
logger.warning("projections compile failed: %s", e)
|
|
|
|
if targets:
|
|
logger.info("VoxCPM2: torch.compile applied to: %s", ", ".join(targets))
|
|
|
|
def _dit_proj_fn(self, lm_h: torch.Tensor, res_h: torch.Tensor) -> torch.Tensor:
|
|
tts = self.tts
|
|
return torch.cat([tts.lm_to_dit_proj(lm_h), tts.res_to_dit_proj(res_h)], dim=-1)
|
|
|
|
def _stop_fn(self, lm_h: torch.Tensor) -> torch.Tensor:
|
|
tts = self.tts
|
|
return tts.stop_head(tts.stop_actn(tts.stop_proj(lm_h)))
|
|
|
|
def _get_cuda_graph_pool(self) -> tuple:
|
|
if self._cuda_graph_pool is None:
|
|
self._cuda_graph_pool = torch.cuda.graph_pool_handle()
|
|
return self._cuda_graph_pool
|
|
|
|
@staticmethod
|
|
def _nullify_volatile_metadata(ctx: Any) -> Any:
|
|
"""Set ``scheduler_metadata`` to None on all attention layers.
|
|
|
|
This is the only tensor FA3 reallocates each step (variable shape).
|
|
All other metadata tensors are persistent model-runner buffers.
|
|
Setting it to None makes FA3 use default scheduling (~0.1ms cost).
|
|
"""
|
|
if not isinstance(ctx.attn_metadata, dict):
|
|
return ctx
|
|
|
|
ctx = copy.copy(ctx)
|
|
new_meta: dict[str, Any] = {}
|
|
for layer_name, meta in ctx.attn_metadata.items():
|
|
if getattr(meta, "scheduler_metadata", None) is not None:
|
|
meta = copy.copy(meta)
|
|
meta.scheduler_metadata = None
|
|
new_meta[layer_name] = meta
|
|
ctx.attn_metadata = new_meta
|
|
return ctx
|
|
|
|
def _capture_graph(
|
|
self,
|
|
model: nn.Module,
|
|
batch_size: int,
|
|
label: str,
|
|
is_residual: bool = False,
|
|
) -> _CapturedGraph:
|
|
"""Capture a CUDA Graph for *model* at *batch_size*."""
|
|
hidden_size = self.config.hidden_size
|
|
dtype = self._side_dtype
|
|
dev = torch.device(self._device)
|
|
pool = self._get_cuda_graph_pool()
|
|
|
|
model.precompute_fused_qkv()
|
|
|
|
g = _CapturedGraph(
|
|
graph=torch.cuda.CUDAGraph(),
|
|
input_embeds=torch.zeros(batch_size, hidden_size, device=dev, dtype=dtype),
|
|
positions=torch.zeros(batch_size, device=dev, dtype=torch.long),
|
|
output=torch.zeros(batch_size, hidden_size, device=dev, dtype=dtype),
|
|
)
|
|
|
|
if is_residual:
|
|
call_kwargs = dict(positions=g.positions, inputs_embeds=g.input_embeds)
|
|
else:
|
|
call_kwargs = dict(input_ids=None, positions=g.positions, inputs_embeds=g.input_embeds)
|
|
|
|
ctx = get_forward_context()
|
|
patched_ctx = self._nullify_volatile_metadata(ctx)
|
|
|
|
with override_forward_context(patched_ctx):
|
|
for _ in range(3):
|
|
_ = model(**call_kwargs)
|
|
|
|
with torch.cuda.graph(g.graph, pool=pool):
|
|
g.output = model(**call_kwargs)
|
|
|
|
logger.info("CUDA Graph captured for %s (batch_size=%d)", label, batch_size)
|
|
return g
|
|
|
|
def _replay_graph(
|
|
self,
|
|
g: _CapturedGraph,
|
|
inputs_embeds: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
batch_size: int,
|
|
) -> torch.Tensor:
|
|
"""Copy fresh inputs into static buffers, then replay.
|
|
|
|
No metadata copy needed: persistent buffers (seq_lens, slot_mapping,
|
|
etc.) are updated in-place by the model runner. scheduler_metadata
|
|
was nullified at capture time so no kernel references it.
|
|
"""
|
|
g.input_embeds[:batch_size].copy_(inputs_embeds[:batch_size])
|
|
g.positions[:batch_size].copy_(positions[:batch_size])
|
|
g.graph.replay()
|
|
return g.output[:batch_size].clone()
|
|
|
|
# -------------------- vllm hooks --------------------
|
|
|
|
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
|
|
return self.model.embed_input_ids(input_ids)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
**kwargs: Any,
|
|
) -> torch.Tensor | IntermediateTensors:
|
|
self._perf.start("forward_total")
|
|
dev = input_ids.device
|
|
|
|
num_reqs = len(self._pending_requests)
|
|
num_decode = sum(1 for _, is_p, _, n in self._pending_requests if not is_p and n == 1)
|
|
is_all_decode = num_decode == num_reqs and num_reqs > 0
|
|
|
|
tts_compiled = getattr(self.tts.feat_decoder.estimator, "_compiled", False) if self._tts is not None else False
|
|
graph_ready = tts_compiled and self._cuda_graph_warmup_steps >= self._cuda_graph_warmup_threshold
|
|
if num_decode > 0:
|
|
self._cuda_graph_warmup_steps += 1
|
|
|
|
can_use_graph = (
|
|
self._enable_cuda_graph and graph_ready and intermediate_tensors is None and inputs_embeds is not None
|
|
)
|
|
|
|
if can_use_graph and is_all_decode and num_reqs <= self._max_cached_graphs:
|
|
self._perf.start("scaffold_fwd")
|
|
if num_reqs not in self._scaffold_graphs:
|
|
self._scaffold_graphs[num_reqs] = self._capture_graph(self.model, num_reqs, "scaffold")
|
|
scaffold_hidden = self._replay_graph(self._scaffold_graphs[num_reqs], inputs_embeds, positions, num_reqs)
|
|
self._perf.stop("scaffold_fwd")
|
|
|
|
else:
|
|
self._perf.start("scaffold_fwd")
|
|
model_output = self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
|
|
self._perf.stop("scaffold_fwd")
|
|
if isinstance(model_output, IntermediateTensors):
|
|
return model_output
|
|
scaffold_hidden = model_output
|
|
if isinstance(scaffold_hidden, tuple):
|
|
scaffold_hidden = scaffold_hidden[0]
|
|
|
|
# Phase 1: per-request FSQ + residual input
|
|
token_offset = 0
|
|
residual_inputs: list[torch.Tensor] = []
|
|
residual_positions: list[torch.Tensor] = []
|
|
req_metas: list[tuple] = []
|
|
|
|
for req_id, is_prefill, _req_embeds, n in self._pending_requests:
|
|
state = self._switch_to_request(req_id)
|
|
req_hidden = scaffold_hidden[token_offset : token_offset + n]
|
|
req_pos = positions[token_offset : token_offset + n]
|
|
|
|
if is_prefill:
|
|
res_input, meta = self._prepare_residual_prefill(state, req_hidden, dev)
|
|
elif state.prefill_completed:
|
|
res_input, meta = self._prepare_residual_decode(state, req_hidden, dev)
|
|
else:
|
|
token_offset += n
|
|
self._results_queue.append((req_id, None))
|
|
self._audio_queue.append((req_id, None))
|
|
continue
|
|
|
|
residual_inputs.append(res_input)
|
|
residual_positions.append(req_pos)
|
|
req_metas.append((state, is_prefill, meta))
|
|
token_offset += n
|
|
|
|
# Phase 2: batch residual_model forward
|
|
if residual_inputs:
|
|
batch_in = torch.cat(residual_inputs, dim=0)
|
|
batch_pos = torch.cat(residual_positions, dim=0)
|
|
|
|
residual_batch_size = batch_in.shape[0]
|
|
use_residual_graph = (
|
|
self._enable_cuda_graph
|
|
and is_all_decode
|
|
and graph_ready
|
|
and residual_batch_size == num_reqs # 1 token per request
|
|
and residual_batch_size <= self._max_cached_graphs
|
|
)
|
|
|
|
self._perf.start("residual_fwd")
|
|
if use_residual_graph:
|
|
if residual_batch_size not in self._residual_graphs:
|
|
self._residual_graphs[residual_batch_size] = self._capture_graph(
|
|
self.residual_model, residual_batch_size, "residual", is_residual=True
|
|
)
|
|
batch_out = self._replay_graph(
|
|
self._residual_graphs[residual_batch_size], batch_in, batch_pos, residual_batch_size
|
|
)
|
|
else:
|
|
batch_out = self.residual_model(batch_pos, batch_in)
|
|
self._perf.stop("residual_fwd")
|
|
|
|
# Phase 3: per-request LocDiT + update
|
|
offset = 0
|
|
for idx, (state, is_prefill, meta) in enumerate(req_metas):
|
|
n = residual_inputs[idx].shape[0]
|
|
res_out = batch_out[offset : offset + n]
|
|
offset += n
|
|
|
|
if is_prefill:
|
|
self._finish_prefill(state, meta, res_out, dev)
|
|
else:
|
|
self._finish_decode(state, meta, res_out, dev)
|
|
|
|
self._results_queue.append((state.request_id, state.precomputed_stop_logits))
|
|
self._audio_queue.append((state.request_id, self._collect_audio(state)))
|
|
|
|
self._pending_requests.clear()
|
|
self._flush_deferred_cleanup()
|
|
self._perf.stop("forward_total")
|
|
return scaffold_hidden
|
|
|
|
# -------------------- prefill / decode helpers --------------------
|
|
|
|
def _prepare_residual_prefill(self, state: _RequestState, base_lm_out: torch.Tensor, dev: Any):
|
|
tts = self.tts
|
|
text_mask, feat_mask, feat, feat_embed = state.prefill_masks
|
|
state.prefill_masks = None
|
|
|
|
tts_len = text_mask.shape[1]
|
|
scaffold_len = base_lm_out.shape[0]
|
|
assert scaffold_len == tts_len, (
|
|
f"voxcpm2 prefill length mismatch: scaffold_len={scaffold_len} tts_len={tts_len}; "
|
|
"caller must pad prompt_token_ids to the full prefill length "
|
|
"(see serving_speech._build_voxcpm2_prompt or the offline example)."
|
|
)
|
|
enc_out = base_lm_out.unsqueeze(0)
|
|
|
|
prefix_feat_cond = (
|
|
feat[:, -1, ...]
|
|
if feat.shape[1] > 0
|
|
else torch.zeros(1, self._patch_size, self._feat_dim, device=dev, dtype=self._side_dtype)
|
|
)
|
|
enc_outputs = tts.fsq_layer(enc_out) * feat_mask.unsqueeze(-1) + enc_out * text_mask.unsqueeze(-1)
|
|
lm_hidden = enc_outputs[:, -1, :]
|
|
|
|
residual_input = tts.fusion_concat_proj(torch.cat([enc_outputs, feat_mask.unsqueeze(-1) * feat_embed], dim=-1))
|
|
meta = {"lm_hidden": lm_hidden, "prefix_feat_cond": prefix_feat_cond}
|
|
return residual_input.squeeze(0), meta
|
|
|
|
def _prepare_residual_decode(self, state: _RequestState, base_lm_out: torch.Tensor, dev: Any):
|
|
tts = self.tts
|
|
state.decode_step_count += 1
|
|
|
|
if state.decode_step_count >= self._max_decode_steps:
|
|
logger.warning("MAX_DECODE_STEPS for %s (%d), forcing stop", state.request_id, state.decode_step_count)
|
|
state.is_stopping = True
|
|
|
|
h = base_lm_out.unsqueeze(0) if base_lm_out.ndim == 1 else base_lm_out
|
|
lm_h = tts.fsq_layer(h)
|
|
if lm_h.ndim == 1:
|
|
lm_h = lm_h.unsqueeze(0)
|
|
|
|
prev = state.prev_feat_embed.to(self._side_dtype)
|
|
if prev.ndim == 1:
|
|
prev = prev.unsqueeze(0)
|
|
res_input = tts.fusion_concat_proj(torch.cat([lm_h, prev], dim=-1))
|
|
return res_input, {"new_lm_hidden": lm_h}
|
|
|
|
def _run_cfm(self, dit_h: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
|
|
if self._cfm_buffers is not None:
|
|
return _optimized_solve_euler(
|
|
self.tts.feat_decoder,
|
|
dit_h,
|
|
self._patch_size,
|
|
cond,
|
|
self._inference_timesteps,
|
|
self._cfg_value,
|
|
self._cfm_buffers,
|
|
cfg_cutoff_ratio=self._cfg_cutoff_ratio,
|
|
perf=self._perf,
|
|
).transpose(1, 2)
|
|
return self.tts.feat_decoder(
|
|
mu=dit_h,
|
|
patch_size=self._patch_size,
|
|
cond=cond,
|
|
n_timesteps=self._inference_timesteps,
|
|
cfg_value=self._cfg_value,
|
|
).transpose(1, 2)
|
|
|
|
def _finish_prefill(self, state: _RequestState, meta: dict, res_out: torch.Tensor, dev: Any):
|
|
tts = self.tts
|
|
lm_hidden = meta["lm_hidden"]
|
|
prefix_feat_cond = meta["prefix_feat_cond"]
|
|
residual_hidden = res_out[-1:, :]
|
|
|
|
state.precomputed_stop_logits = tts.stop_head(tts.stop_actn(tts.stop_proj(lm_hidden))).detach()
|
|
dit_h = torch.cat([tts.lm_to_dit_proj(lm_hidden), tts.res_to_dit_proj(residual_hidden)], dim=-1)
|
|
|
|
self._setup_cfm_buffers()
|
|
if self._enable_torch_compile:
|
|
self._setup_torch_compile()
|
|
|
|
pred_feat = self._run_cfm(dit_h, prefix_feat_cond.transpose(1, 2).contiguous())
|
|
|
|
with torch.no_grad():
|
|
curr_embed = tts.enc_to_lm_proj(tts.feat_encoder(pred_feat.unsqueeze(1))).squeeze(1)
|
|
|
|
state.curr_embed_for_next = curr_embed.detach()
|
|
state.prev_feat_embed = curr_embed.detach()
|
|
state.curr_prefix_feat_cond = pred_feat[0].detach()
|
|
state.last_audio_patch_gpu = pred_feat.detach()
|
|
state.decode_step_count = 0
|
|
state.request_start_time = time.perf_counter()
|
|
state.prefill_completed = True
|
|
|
|
if logger.isEnabledFor(logging.DEBUG):
|
|
# Only compute the norm (which forces a GPU->CPU sync) if we will log it.
|
|
logger.debug("PREFILL[%s]: patch norm=%.4f", state.request_id, pred_feat.norm().item())
|
|
self._perf.reset()
|
|
|
|
def _finish_decode(self, state: _RequestState, meta: dict, res_out: torch.Tensor, dev: Any):
|
|
self._perf.start("decode_step")
|
|
tts = self.tts
|
|
|
|
lm_h = meta["new_lm_hidden"]
|
|
res_h = res_out.unsqueeze(0) if res_out.ndim == 1 else res_out
|
|
|
|
dit_proj = getattr(self, "_compiled_dit_proj", None) or self._dit_proj_fn
|
|
stop_fn = getattr(self, "_compiled_stop_fn", None) or self._stop_fn
|
|
|
|
dit_h = dit_proj(lm_h, res_h)
|
|
pfc = state.curr_prefix_feat_cond.to(self._side_dtype)
|
|
if pfc.ndim == 2:
|
|
pfc = pfc.unsqueeze(0)
|
|
|
|
pred_feat = self._run_cfm(dit_h, pfc.transpose(1, 2).contiguous())
|
|
next_embed = tts.enc_to_lm_proj(tts.feat_encoder(pred_feat.unsqueeze(1))).squeeze(1)
|
|
|
|
state.precomputed_stop_logits = stop_fn(lm_h).detach()
|
|
state.curr_embed_for_next = next_embed.detach()
|
|
state.prev_feat_embed = next_embed.detach()
|
|
state.curr_prefix_feat_cond = pred_feat[0].detach()
|
|
state.last_audio_patch_gpu = pred_feat.detach()
|
|
|
|
self._perf.stop("decode_step")
|
|
if _ENABLE_PROFILING and state.decode_step_count % 20 == 0:
|
|
logger.info("Step %d[%s]:\n%s", state.decode_step_count, state.request_id, self._perf.breakdown())
|
|
|
|
# -------------------- audio collection --------------------
|
|
|
|
def _collect_audio(self, state: _RequestState) -> torch.Tensor | None:
|
|
"""Per-step sliding-window VAE decode (nanovllm pattern).
|
|
|
|
Each decode step feeds ``[decode_pad, new_patch]`` through the VAE
|
|
and slices out only the audio region corresponding to the new patch.
|
|
The pad buffer (last ``_n_decode_pad_frames`` latent frames) provides
|
|
the receptive-field context needed by the VAE's transposed convolutions,
|
|
eliminating boundary artifacts between chunks.
|
|
|
|
Returns the delta audio chunk (not cumulative) so the output processor
|
|
can stream each chunk to the client independently.
|
|
"""
|
|
patch = state.last_audio_patch_gpu
|
|
if patch is None:
|
|
return None
|
|
state.last_audio_patch_gpu = None
|
|
|
|
# patch shape: (patch_size, feat_dim) or (1, patch_size, feat_dim)
|
|
new_latent = patch.reshape(-1, self._feat_dim).to(torch.float32)
|
|
n_new = new_latent.shape[0] # = patch_size (typically 4)
|
|
|
|
self._perf.start("vae_decode")
|
|
|
|
# Build VAE input: [pad_frames | new_latent]
|
|
if state.decode_pad is not None:
|
|
vae_input = torch.cat([state.decode_pad, new_latent], dim=0)
|
|
pad_frames = state.decode_pad.shape[0]
|
|
else:
|
|
vae_input = new_latent
|
|
pad_frames = 0
|
|
|
|
# VAE decode: (1, feat_dim, T_frames) -> (1, 1, T_samples)
|
|
feat = vae_input.unsqueeze(0).transpose(1, 2).contiguous()
|
|
with torch.no_grad():
|
|
audio = self.tts.audio_vae.decode(feat.to(self._device)).reshape(-1)
|
|
|
|
# Slice out only the new audio (after the pad region).
|
|
# Each latent frame maps to decoder_chunk_size audio samples.
|
|
dcs = int(getattr(self.tts.audio_vae, "decode_chunk_size", audio.numel() // vae_input.shape[0]))
|
|
new_audio = audio[pad_frames * dcs : (pad_frames + n_new) * dcs].detach().cpu().float()
|
|
|
|
# Roll the pad buffer: keep last N latent frames as context for next step.
|
|
all_latents = vae_input # [pad + new]
|
|
state.decode_pad = all_latents[-self._n_decode_pad_frames :].detach()
|
|
|
|
self._perf.stop("vae_decode")
|
|
return new_audio
|
|
|
|
# -------------------- compute_logits --------------------
|
|
|
|
def compute_logits(
|
|
self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None
|
|
) -> torch.Tensor | None:
|
|
if isinstance(hidden_states, OmniOutput):
|
|
hidden_states = hidden_states.text_hidden_states
|
|
if hidden_states is None:
|
|
return None
|
|
|
|
bsz = hidden_states.shape[0]
|
|
logits = torch.full(
|
|
(bsz, self.config.vocab_size), float("-inf"), device=hidden_states.device, dtype=hidden_states.dtype
|
|
)
|
|
|
|
if self._results_queue:
|
|
for i, (req_id, stop_logits) in enumerate(self._results_queue):
|
|
if i >= bsz:
|
|
break
|
|
state = self._active_states.get(req_id)
|
|
if stop_logits is not None:
|
|
if state is not None and state.is_stopping:
|
|
logits[i, 0] = 0.0
|
|
logits[i, 1] = 1.0
|
|
state.precomputed_stop_logits = None
|
|
else:
|
|
logits[i, 0] = stop_logits[0, 0]
|
|
logits[i, 1] = stop_logits[0, 1]
|
|
if state is not None:
|
|
state.is_stopping = bool(stop_logits[0, 1] > stop_logits[0, 0])
|
|
state.precomputed_stop_logits = None
|
|
elif state and state.prefill_completed:
|
|
logits[i, 1] = 1.0
|
|
else:
|
|
logits[i, 0] = 1.0
|
|
self._results_queue.clear()
|
|
else:
|
|
logits[:, 0] = 1.0
|
|
return logits
|
|
|
|
# -------------------- omni output --------------------
|
|
|
|
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
|
|
if isinstance(model_outputs, OmniOutput):
|
|
return model_outputs
|
|
|
|
mm: dict[str, Any] = {}
|
|
if self._audio_queue:
|
|
audio_by_req = {rid: audio for rid, audio in self._audio_queue}
|
|
order = [r for r, _ in self._audio_queue]
|
|
mm["model_outputs"] = [audio_by_req.get(r) for r in order]
|
|
mm["sr"] = [torch.tensor(self._sample_rate, dtype=torch.int32) for _ in order]
|
|
self._audio_queue.clear()
|
|
|
|
return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm)
|
|
|
|
# -------------------- Chinese token splitting --------------------
|
|
|
|
def _get_multichar_zh_split(self) -> dict[int, list[int]]:
|
|
"""Lazy-build {multichar_chinese_token_id: [char_id, ...]} map."""
|
|
if self._multichar_zh_split is not None:
|
|
return self._multichar_zh_split
|
|
base_tokenizer = self.tts.text_tokenizer.tokenizer
|
|
self._multichar_zh_split = build_cjk_split_map(base_tokenizer)
|
|
logger.info("VoxCPM2: built multichar Chinese split map (%d entries)", len(self._multichar_zh_split))
|
|
return self._multichar_zh_split
|
|
|
|
# -------------------- preprocess / postprocess --------------------
|
|
|
|
def preprocess(
|
|
self, input_ids: torch.Tensor, input_embeds: torch.Tensor | None, **info_dict: Any
|
|
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
|
|
additional = info_dict.get("additional_information")
|
|
if isinstance(additional, dict):
|
|
merged = {k: v for k, v in info_dict.items() if k != "additional_information"}
|
|
for k, v in additional.items():
|
|
merged.setdefault(k, v)
|
|
info_dict = merged
|
|
|
|
span_len = int(input_ids.shape[0])
|
|
dev = input_ids.device
|
|
req_id = info_dict.get("request_id", "default")
|
|
is_prefill = span_len > 1
|
|
|
|
if is_prefill:
|
|
# Do not evict state here: _pending_requests is a per-step prefix,
|
|
# not the full batch. Cleanup is driven by on_requests_finished ->
|
|
# _flush_deferred_cleanup (fed by vLLM scheduler._free_request via
|
|
# gpu_ar_model_runner.py).
|
|
real = info_dict.get("text_token_ids")
|
|
token_ids = input_ids.tolist() if real is None else real[0]
|
|
# Fail-fast: unsplit multichar Chinese IDs in input_ids means the
|
|
# serving layer didn't pre-split. Silent fixup here would cause
|
|
# input_ids/embeds length mismatch (scheduler slot count is fixed).
|
|
split_map = self._get_multichar_zh_split()
|
|
if split_map and any(tid in split_map for tid in token_ids):
|
|
raise ValueError(
|
|
"VoxCPM2 preprocess received unsplit multichar Chinese "
|
|
"token IDs. The serving layer must send prompt_token_ids "
|
|
"with single-char CJK IDs (see _voxcpm2_encode)."
|
|
)
|
|
if token_ids and token_ids[0] == self.config.bos_token_id:
|
|
token_ids = token_ids[1:]
|
|
|
|
state = self._get_or_create_state(req_id)
|
|
state.prefill_text = ""
|
|
state.decode_pad = None
|
|
state.prefill_completed = False
|
|
state.decode_step_count = 0
|
|
state.precomputed_stop_logits = None
|
|
state.last_audio_patch_gpu = None
|
|
state.curr_embed_for_next = None
|
|
state.prev_feat_embed = None
|
|
state.curr_prefix_feat_cond = None
|
|
state.is_stopping = False
|
|
|
|
# Voice clone / continuation
|
|
ref_audio = info_dict.get("reference_audio") or info_dict.get("ref_audio")
|
|
prompt_audio = info_dict.get("prompt_audio")
|
|
prompt_text = info_dict.get("prompt_text")
|
|
if isinstance(ref_audio, list):
|
|
ref_audio = ref_audio[0] if ref_audio else None
|
|
if isinstance(prompt_audio, list):
|
|
prompt_audio = prompt_audio[0] if prompt_audio else None
|
|
if isinstance(prompt_text, list):
|
|
prompt_text = prompt_text[0] if prompt_text else None
|
|
|
|
state.prompt_cache = None
|
|
voice_name = info_dict.get("voice_name")
|
|
if isinstance(voice_name, list):
|
|
voice_name = voice_name[0] if voice_name else None
|
|
_created_at = int(info_dict.get("voice_created_at") or 0)
|
|
|
|
if ref_audio or (prompt_audio and prompt_text):
|
|
# Check speaker cache for reference-only mode
|
|
if voice_name and ref_audio and not prompt_audio:
|
|
_cache_key = self._speaker_cache.make_cache_key(
|
|
voice_name, model_type="voxcpm2", created_at=_created_at
|
|
)
|
|
cached = self._speaker_cache.get(_cache_key)
|
|
if cached is not None:
|
|
state.prompt_cache = {
|
|
"mode": "reference",
|
|
"ref_audio_feat": cached["ref_audio_feat"].clone(),
|
|
}
|
|
logger.debug("Speaker cache HIT for VoxCPM2 speaker '%s'", voice_name)
|
|
|
|
if state.prompt_cache is None:
|
|
try:
|
|
state.prompt_cache = self._build_prompt_cache(
|
|
ref_audio=ref_audio,
|
|
prompt_audio=prompt_audio,
|
|
prompt_text=prompt_text,
|
|
)
|
|
if (
|
|
voice_name
|
|
and state.prompt_cache is not None
|
|
and state.prompt_cache.get("mode") == "reference"
|
|
and "ref_audio_feat" in state.prompt_cache
|
|
):
|
|
_key = self._speaker_cache.make_cache_key(
|
|
voice_name, model_type="voxcpm2", created_at=_created_at
|
|
)
|
|
self._speaker_cache.put(
|
|
_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()}
|
|
)
|
|
logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name)
|
|
except Exception as e:
|
|
logger.warning("build_prompt_cache failed: %s; falling back to zero-shot", e)
|
|
state.prompt_cache = None
|
|
|
|
inputs = self._build_prefill_inputs(token_ids, dev, req_id)
|
|
tts = self.tts
|
|
feat_embed = tts.enc_to_lm_proj(tts.feat_encoder(inputs["audio_feat"]))
|
|
text_embed = self.model.embed_input_ids(inputs["text_token"].to(dev))
|
|
text_mask, feat_mask = inputs["text_mask"], inputs["audio_mask"]
|
|
embeds = (text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed).squeeze(0)
|
|
state.prefill_masks = (text_mask, feat_mask, inputs["audio_feat"], feat_embed)
|
|
else:
|
|
state = self._active_states.get(req_id)
|
|
curr = state.curr_embed_for_next if state else None
|
|
if curr is not None:
|
|
embeds = curr.to(dev, dtype=self._side_dtype).reshape(1, -1)
|
|
else:
|
|
embeds = torch.zeros(1, self.config.hidden_size, device=dev, dtype=self._side_dtype)
|
|
|
|
self._pending_requests.append((req_id, is_prefill, embeds, span_len))
|
|
return input_ids, embeds, {}
|
|
|
|
def postprocess(self, hidden_states: torch.Tensor, **info: Any) -> dict[str, Any]:
|
|
req_id = info.get("request_id", self._current_request_id or "default")
|
|
if _ENABLE_PROFILING:
|
|
state = self._active_states.get(req_id)
|
|
if state and state.decode_step_count > 0:
|
|
logger.info(
|
|
"REQUEST DONE[%s]: %d steps, %.2fs\n%s",
|
|
req_id,
|
|
state.decode_step_count,
|
|
time.perf_counter() - state.request_start_time,
|
|
self._perf.breakdown(),
|
|
)
|
|
return {}
|
|
|
|
# -------------------- build prefill inputs --------------------
|
|
|
|
def _build_prefill_inputs(self, token_ids: list[int], dev: Any, req_id: str = "default") -> dict:
|
|
tts = self.tts
|
|
dtype = self._side_dtype
|
|
state = self._active_states.get(req_id)
|
|
cache = state.prompt_cache if state else None
|
|
mode = cache.get("mode", "continuation") if cache else "zero_shot"
|
|
|
|
if cache and mode in ("continuation", "ref_continuation"):
|
|
prompt_text = cache.get("prompt_text", "")
|
|
prompt_ids = list(tts.text_tokenizer(prompt_text)) if prompt_text else []
|
|
all_ids = prompt_ids + token_ids
|
|
else:
|
|
all_ids = token_ids
|
|
|
|
text_token = torch.tensor(all_ids, dtype=torch.int32)
|
|
text_token = torch.cat([text_token, torch.tensor([tts.audio_start_token], dtype=torch.int32)], dim=-1)
|
|
text_len = text_token.shape[0]
|
|
latent_dim = tts.audio_vae.latent_dim
|
|
ps = self._patch_size
|
|
|
|
if mode in ("zero_shot", "continuation"):
|
|
audio_feat = cache["audio_feat"] if cache else torch.empty((0, ps, latent_dim), dtype=torch.float32)
|
|
a_len = audio_feat.size(0)
|
|
text_token = torch.cat([text_token, torch.zeros(a_len, dtype=torch.int32)])
|
|
audio_feat = torch.cat([torch.zeros((text_len, ps, latent_dim), dtype=torch.float32), audio_feat])
|
|
text_mask = torch.cat([torch.ones(text_len, dtype=torch.int32), torch.zeros(a_len, dtype=torch.int32)])
|
|
audio_mask = torch.cat([torch.zeros(text_len, dtype=torch.int32), torch.ones(a_len, dtype=torch.int32)])
|
|
elif mode == "reference":
|
|
ref = cache["ref_audio_feat"]
|
|
rt, rf, rtm, ram = tts._make_ref_prefix(ref, text_token.device)
|
|
text_token = torch.cat([rt.cpu(), text_token])
|
|
audio_feat = torch.cat([rf.cpu(), torch.zeros((text_len, ps, latent_dim), dtype=torch.float32)])
|
|
text_mask = torch.cat([rtm.cpu(), torch.ones(text_len, dtype=torch.int32)])
|
|
audio_mask = torch.cat([ram.cpu(), torch.zeros(text_len, dtype=torch.int32)])
|
|
else: # ref_continuation
|
|
ref = cache["ref_audio_feat"]
|
|
prompt = cache["audio_feat"]
|
|
p_len = prompt.size(0)
|
|
rt, rf, rtm, ram = tts._make_ref_prefix(ref, text_token.device)
|
|
text_token = torch.cat([rt.cpu(), text_token, torch.zeros(p_len, dtype=torch.int32)])
|
|
audio_feat = torch.cat([rf.cpu(), torch.zeros((text_len, ps, latent_dim), dtype=torch.float32), prompt])
|
|
ones_t = torch.ones(text_len, dtype=torch.int32)
|
|
zeros_p = torch.zeros(p_len, dtype=torch.int32)
|
|
zeros_t = torch.zeros(text_len, dtype=torch.int32)
|
|
ones_p = torch.ones(p_len, dtype=torch.int32)
|
|
text_mask = torch.cat([rtm.cpu(), ones_t, zeros_p])
|
|
audio_mask = torch.cat([ram.cpu(), zeros_t, ones_p])
|
|
|
|
return {
|
|
"text_token": text_token.unsqueeze(0).to(dev),
|
|
"audio_feat": audio_feat.unsqueeze(0).to(dev).to(dtype),
|
|
"text_mask": text_mask.unsqueeze(0).to(dev),
|
|
"audio_mask": audio_mask.unsqueeze(0).to(dev),
|
|
}
|
|
|
|
# -------------------- weight loading --------------------
|
|
|
|
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"base_lm.": "model."})
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
def _base_lm_only(ws):
|
|
for name, tensor in ws:
|
|
if name.startswith("base_lm."):
|
|
yield name, tensor
|
|
|
|
loader = AutoWeightsLoader(self)
|
|
loaded = loader.load_weights(_base_lm_only(weights), mapper=self.hf_to_vllm_mapper)
|
|
|
|
# _tts and residual_model are constructed and populated eagerly in
|
|
# __init__ via VoxCPM.from_pretrained; here we only need to mark their
|
|
# params as loaded so AutoWeightsLoader's strict check doesn't flag
|
|
# them as missing from the checkpoint.
|
|
loaded |= {name for name, _ in self.named_parameters() if name.startswith(("_tts.", "residual_model."))}
|
|
|
|
logger.info(
|
|
"Loaded VoxCPM2 (patch=%d, feat_dim=%d, dtype=%s)", self._patch_size, self._feat_dim, self._side_dtype
|
|
)
|
|
return loaded
|