From 226cb4a38cee37d2cc6a3d509703b5473f815e81 Mon Sep 17 00:00:00 2001 From: boatman <109857087+sphinxkkkbc@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:44:53 +0800 Subject: [PATCH] [Bugfix][VoxCPM2] Fix Chunked Prefill Length Mismatch (#5416) Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com> --- .../voxcpm2/test_talker_state_eviction.py | 161 ++++++++++++++++++ .../models/voxcpm2/voxcpm2_talker.py | 124 +++++++++----- 2 files changed, 241 insertions(+), 44 deletions(-) diff --git a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py index 3938051d3..e0b52d1aa 100644 --- a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py +++ b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py @@ -235,6 +235,9 @@ class _NoopPerf: def stop(self, name: str) -> None: pass + def reset(self) -> None: + pass + class _FakeTTS: feat_decoder = SimpleNamespace(estimator=SimpleNamespace(_compiled=False)) @@ -253,6 +256,164 @@ class _FakeTTS: return x.squeeze(1).mean(dim=1).repeat(1, 4) +class TestChunkedPrefillContract: + @staticmethod + def _make_talker(prompt_len: int = 30): + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import _PrefillInputs + + class FakeTTS: + def feat_encoder(self, x: torch.Tensor) -> torch.Tensor: + return x + + def enc_to_lm_proj(self, x: torch.Tensor) -> torch.Tensor: + return x.squeeze(2) + + def stop_proj(self, x: torch.Tensor) -> torch.Tensor: + return x + + def stop_actn(self, x: torch.Tensor) -> torch.Tensor: + return x + + def stop_head(self, x: torch.Tensor) -> torch.Tensor: + return x[:, :2] + + def lm_to_dit_proj(self, x: torch.Tensor) -> torch.Tensor: + return x + + def res_to_dit_proj(self, x: torch.Tensor) -> torch.Tensor: + return x + + class FakeScaffold: + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return input_ids.to(torch.float32).unsqueeze(-1).repeat(1, 1, 4) + + talker = _make_bare_talker() + talker._tts = FakeTTS() + talker.model = FakeScaffold() + talker.config = SimpleNamespace(bos_token_id=-1, hidden_size=4) + talker._side_dtype = torch.float32 + talker._get_multichar_zh_split = lambda: {} + talker._perf = _NoopPerf() + talker._enable_torch_compile = False + talker._setup_cfm_buffers = lambda: None + talker._run_cfm_for_state = lambda *_args: torch.ones(1, 1, 4) + + text_token = torch.arange(prompt_len, dtype=torch.int32).unsqueeze(0) + audio_feat = torch.zeros(1, prompt_len, 1, 4) + text_mask = torch.ones(1, prompt_len, dtype=torch.int32) + audio_mask = torch.zeros(1, prompt_len, dtype=torch.int32) + talker._build_prefill_inputs = lambda *_args, **_kwargs: _PrefillInputs( + text_token=text_token, + audio_feat=audio_feat, + text_mask=text_mask, + audio_mask=audio_mask, + ) + return talker + + def test_preprocess_slices_full_embeddings_by_request_position(self) -> None: + talker = self._make_talker() + + _, first, _ = talker.preprocess( + input_ids=torch.ones(15, dtype=torch.int32), + input_embeds=None, + request_id="req", + text_token_ids=[list(range(29))], + _omni_is_prefill=True, + _omni_num_computed_tokens=0, + ) + state = talker._active_states["req"] + state.decode_step_count = 7 + + _, second, _ = talker.preprocess( + input_ids=torch.ones(15, dtype=torch.int32), + input_embeds=None, + request_id="req", + text_token_ids=[list(range(29))], + _omni_is_prefill=True, + _omni_num_computed_tokens=15, + ) + + torch.testing.assert_close(first[:, 0], torch.arange(15, dtype=torch.float32)) + torch.testing.assert_close(second[:, 0], torch.arange(15, 30, dtype=torch.float32)) + assert state.decode_step_count == 7 + assert state.is_last_prefill_chunk is True + + def test_prefill_cache_lifecycle_across_chunks(self) -> None: + talker = self._make_talker() + prompt_cache = {"mode": "reference", "ref_audio_feat": torch.ones(1)} + build_calls = [] + + def _build_prompt_cache(**_kwargs): + build_calls.append(1) + return prompt_cache + + talker._build_prompt_cache = _build_prompt_cache + + talker.preprocess( + input_ids=torch.ones(15, dtype=torch.int32), + input_embeds=None, + request_id="req", + text_token_ids=[list(range(29))], + reference_audio="ref.wav", + _omni_is_prefill=True, + _omni_num_computed_tokens=0, + ) + state = talker._active_states["req"] + prefill_embeds = state.prefill_embeds + prefill_masks = state.prefill_masks + + talker.preprocess( + input_ids=torch.ones(15, dtype=torch.int32), + input_embeds=None, + request_id="req", + text_token_ids=[list(range(29))], + reference_audio="ref.wav", + _omni_is_prefill=True, + _omni_num_computed_tokens=15, + ) + + assert build_calls == [1] + assert state.prompt_cache is prompt_cache + assert state.prefill_embeds is prefill_embeds + assert state.prefill_masks is prefill_masks + + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import _PrefillResidualMeta + + talker._finish_prefill( + state, + _PrefillResidualMeta( + lm_hidden=torch.ones(1, 4), + prefix_feat_cond=torch.ones(1, 1, 4), + ), + res_out=torch.ones(1, 4), + dev=torch.device("cpu"), + ) + + assert state.prefill_completed is True + assert state.prompt_cache is None + assert state.prefill_embeds is None + assert state.prefill_masks is None + assert talker._active_states["req"] is state + + def test_single_token_tail_is_still_prefill(self) -> None: + talker = self._make_talker() + + _, embeds, _ = talker.preprocess( + input_ids=torch.ones(1, dtype=torch.int32), + input_embeds=None, + request_id="req", + text_token_ids=[list(range(29))], + _omni_is_prefill=True, + _omni_num_computed_tokens=29, + ) + is_last_prefill_chunk = talker._active_states["req"].is_last_prefill_chunk + is_prefill = talker._pending_requests[-1][1] + assert embeds.shape == (1, 4) + assert embeds[0, 0] == 29 + assert is_prefill is True + assert is_last_prefill_chunk is True + + class TestDecodeBatchContract: def test_forward_skip_pending_prefill_keeps_audio_queue_aligned(self) -> None: class FakeScaffold: diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 94ac454d8..3a1677d49 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -264,10 +264,12 @@ class _RequestState: decode_step_count: int = 0 request_start_time: float = 0.0 prefill_completed: bool = False + is_last_prefill_chunk: bool = False prompt_cache: dict | None = None prefill_masks: tuple | None = None is_stopping: bool = False precomputed_is_stopping: bool | None = None + prefill_embeds: torch.Tensor | None = None @dataclasses.dataclass @@ -2054,7 +2056,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if is_prefill: flush_decode_fsq_batch() - res_input, meta = self._prepare_residual_prefill(state, req_hidden, dev) + res_input, meta = self._prepare_residual_prefill(state, req_hidden, dev, req_pos) elif state.prefill_completed: if ( self._enable_batched_fsq_fusion @@ -2141,7 +2143,8 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): offset += n if is_prefill: - prefill_batch.append((state, meta, res_out)) + if state.is_last_prefill_chunk: + prefill_batch.append((state, meta, res_out)) else: self._finish_decode(state, meta, res_out) decoded_states.append(state) @@ -2152,20 +2155,28 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): for state, meta, res_out in prefill_batch: self._finish_prefill(state, meta, res_out, dev) - collect_states = [state for state, _, _ in req_metas] - self._precompute_stop_flags_for_audio_collect(collect_states) - ready_audio_by_req = self._drain_ready_audio_copies_for_states(collect_states) - audio_by_req = self._collect_audio_batch( - collect_states, - initial_delayed_chunks_by_req={ - state.request_id: ready_audio_by_req.get(state.request_id) - for state, is_prefill, _ in req_metas - if not is_prefill - }, - ) + collect_states = [ + state for state, is_prefill, _ in req_metas if not is_prefill or state.is_last_prefill_chunk + ] + audio_by_req: dict[str, torch.Tensor | None] = {} + if collect_states: + self._precompute_stop_flags_for_audio_collect(collect_states) + ready_audio_by_req = self._drain_ready_audio_copies_for_states(collect_states) + audio_by_req = self._collect_audio_batch( + collect_states, + initial_delayed_chunks_by_req={ + state.request_id: ready_audio_by_req.get(state.request_id) + for state, is_prefill, _ in req_metas + if not is_prefill + }, + ) for state, is_prefill, _ in req_metas: - self._results_queue.append((state.request_id, state.precomputed_stop_logits)) - self._audio_queue.append((state.request_id, audio_by_req.get(state.request_id))) + if is_prefill and not state.is_last_prefill_chunk: + self._results_queue.append((state.request_id, None)) + self._audio_queue.append((state.request_id, None)) + else: + self._results_queue.append((state.request_id, state.precomputed_stop_logits)) + self._audio_queue.append((state.request_id, audio_by_req.get(state.request_id))) self._pending_requests.clear() self._flush_deferred_cleanup() @@ -2179,11 +2190,14 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): state: _RequestState, base_lm_out: torch.Tensor, dev: torch.device, + req_pos: torch.Tensor, ) -> tuple[torch.Tensor, _PrefillResidualMeta]: tts = self.tts text_mask, feat_mask, feat, feat_embed = state.prefill_masks - state.prefill_masks = None - + text_mask = text_mask[:, req_pos] + feat_mask = feat_mask[:, req_pos] + feat = feat[:, req_pos] + feat_embed = feat_embed[:, req_pos] tts_len = text_mask.shape[1] scaffold_len = base_lm_out.shape[0] assert scaffold_len == tts_len, ( @@ -2424,6 +2438,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): state.decode_step_count = 0 state.request_start_time = time.perf_counter() state.prefill_completed = True + state.prefill_masks = None + state.prefill_embeds = None + state.prompt_cache = None self._perf.stop("prefill_tail_batch") if self._enable_profiling: @@ -2465,6 +2482,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): state.decode_step_count = 0 state.request_start_time = time.perf_counter() state.prefill_completed = True + state.prefill_embeds = None + state.prefill_masks = None + state.prompt_cache = None self._perf.stop("prefill_tail") if logger.isEnabledFor(logging.DEBUG): @@ -3033,7 +3053,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): span_len = int(input_ids.shape[0]) dev = input_ids.device req_id = info_dict.get("request_id", "default") - is_prefill = span_len > 1 + is_prefill = bool(info_dict.get("_omni_is_prefill", span_len > 1)) if is_prefill: # Do not evict state here: _pending_requests is a per-step prefix, @@ -3056,22 +3076,29 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): token_ids = token_ids[1:] state = self._get_or_create_state(req_id) - state.decode_pad = None - state.prefill_completed = False - state.decode_step_count = 0 - state.precomputed_stop_logits = None - state.precomputed_is_stopping = None - state.last_audio_patch_gpu = None - if not hasattr(state, "pending_audio_chunks_gpu"): - state.pending_audio_chunks_gpu = [] - if not hasattr(state, "pending_audio_copies"): - state.pending_audio_copies = [] - state.pending_audio_chunks_gpu.clear() - state.pending_audio_copies.clear() - state.curr_embed_for_next = None - state.prev_feat_embed = None - state.curr_prefix_feat_cond = None - state.is_stopping = False + num_computed_tokens = int(info_dict.get("_omni_num_computed_tokens", 0)) + is_first_prefill_chunk = num_computed_tokens == 0 + if is_first_prefill_chunk: + state.decode_pad = None + state.prefill_completed = False + state.is_last_prefill_chunk = False + state.decode_step_count = 0 + state.precomputed_stop_logits = None + state.precomputed_is_stopping = None + state.last_audio_patch_gpu = None + if not hasattr(state, "pending_audio_chunks_gpu"): + state.pending_audio_chunks_gpu = [] + if not hasattr(state, "pending_audio_copies"): + state.pending_audio_copies = [] + state.pending_audio_chunks_gpu.clear() + state.pending_audio_copies.clear() + state.curr_embed_for_next = None + state.prev_feat_embed = None + state.curr_prefix_feat_cond = None + state.is_stopping = False + state.prompt_cache = None + state.prefill_embeds = None + state.prefill_masks = None # Voice clone / continuation ref_audio = info_dict.get("reference_audio") or info_dict.get("ref_audio") @@ -3088,13 +3115,12 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): voice_profile = voice_profile[0] if voice_profile else None requires_precomputed_cache = isinstance(voice_profile, dict) - 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 voice_name: + if state.prompt_cache is None and voice_name: _cache_key = self._speaker_cache.make_cache_key( voice_name, model_type="voxcpm2", created_at=_created_at ) @@ -3131,14 +3157,24 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): _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) - - 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) + if state.prefill_embeds is None or state.prefill_masks is 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) + state.prefill_embeds = embeds + chunk_end = num_computed_tokens + span_len + state.is_last_prefill_chunk = chunk_end == state.prefill_embeds.shape[0] + if chunk_end > state.prefill_embeds.shape[0]: + raise ValueError( + "VoxCPM2 chunked prefill exceeds the constructed prompt: " + f"request_id={req_id} start={num_computed_tokens} end={chunk_end} " + f"prompt_len={state.prefill_embeds.shape[0]}" + ) + embeds = state.prefill_embeds[num_computed_tokens:chunk_end] else: state = self._active_states.get(req_id) curr = state.curr_embed_for_next if state else None