diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index cacc8142f4c..91c17c5037d 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for StageConfigFactory and related classes. """ @@ -1661,6 +1661,7 @@ def test_deep_merge_does_not_mutate_inputs(self): ] } original_foo = fake_config["stages"][0]["foo"] + assert isinstance(original_foo, dict) original_b = original_foo["b"] with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): @@ -2128,7 +2129,12 @@ def test_minicpmo_overlays_inherit_talker_sampling_params(self): sampling = deploy.stages[1].default_sampling_params assert sampling is not None, f"{filename} stage 1 lost default_sampling_params" assert sampling["temperature"] == 0.8, filename - assert sampling["top_k"] == 100, filename + # Upstream utils.TTSSamplingParams. min_tokens is not checked here: + # the duplex overlay drops it to 0 because the model budgets each + # generate_chunk itself. + assert sampling["top_k"] == 25, filename + assert sampling["top_p"] == 0.85, filename + assert sampling["repetition_penalty"] == 1.05, filename def test_ci_inherits_from_main(self): ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml")) @@ -2828,13 +2834,14 @@ def pull_files(self, model_path, allow_pattern=None, ignore_pattern=None): monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel) class Recorder: + def __init__(self) -> None: + self.pulls = pulls + def set_files(self, files: list[tuple[str, str]]) -> None: materialized_files.clear() materialized_files.extend(files) - recorder = Recorder() - recorder.pulls = pulls - return recorder + return Recorder() def test_passthrough_for_non_uri(self, fake_object_storage): assert _materialize_object_storage_configs("org/model") == "org/model" diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py index 5a21e8692c3..e7fa50c1611 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """E2E online expansion tests for MiniCPM-o 4.5. These cover modality combinations, separate Code2Wav behavior, sequential @@ -193,7 +196,7 @@ def test_text_to_audio_long_form_001(omni_server, openai_client) -> None: """ messages = dummy_messages_from_mix_data( system_prompt=get_system_prompt(), - content_text="帮我讲一个300字的故事.", + content_text="帮我讲一个100字的故事.", ) request_config = { @@ -202,10 +205,7 @@ def test_text_to_audio_long_form_001(omni_server, openai_client) -> None: "stream": True, "extra_body": _TTS_EXTRA_BODY, } - responses = openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) - text = responses[0].text_content if responses else "" - word_count = len(text.split()) - assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}" + openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py index 4115c1a291c..5a3e7d67d73 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Request-alignment tests for MiniCPM-o 4.5's native Talker.""" from __future__ import annotations from dataclasses import dataclass +from typing import Any import pytest import torch @@ -17,7 +18,9 @@ ConditionalChatTTSConfig, ) from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + _CODEC_PENALTY_WINDOW, _DUPLEX_CODEC_TOKENS_PER_CHUNK, + _OFFLINE_CODEC_MAX_NEW_TOKENS, _REPETITION_PENALTY_CHUNK_SIZE, MiniCPMO45OmniTTSForConditionalGeneration, _apply_batched_repetition_penalty, @@ -37,14 +40,30 @@ class _PenaltyMetadata: prompt_token_ids: torch.Tensor | None +@dataclass +class _SamplerOutput: + """Stand-in for vLLM's SamplerOutput.""" + + sampled_token_ids: torch.Tensor + + +@dataclass +class _CodecSamplingMetadata: + """Stand-in for the codec-penalty slice of vLLM's SamplingMetadata.""" + + repetition_penalties: torch.Tensor + prompt_token_ids: torch.Tensor | None = None + no_penalties: bool = False + + class _FakeNativeTalker(nn.Module): has_preprocess = True def __init__(self) -> None: super().__init__() - self.forward_kwargs = None + self.forward_kwargs: dict[str, Any] | None = None - def forward(self, **kwargs): + def forward(self, **kwargs: Any) -> torch.Tensor: self.forward_kwargs = kwargs return torch.ones(2, 4) @@ -59,7 +78,8 @@ def test_wrapper_always_delegates_talker_to_native_ar_path() -> None: model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) nn.Module.__init__(model) model.model_stage = "tts" - model.talker = _FakeNativeTalker() + talker = _FakeNativeTalker() + model.talker = talker output = model( input_ids=torch.tensor([1, 2]), @@ -68,7 +88,8 @@ def test_wrapper_always_delegates_talker_to_native_ar_path() -> None: ) assert output.shape == (2, 4) - assert model.talker.forward_kwargs["model_intermediate_buffer"][0]["request_id"] == "req" + assert talker.forward_kwargs is not None + assert talker.forward_kwargs["model_intermediate_buffer"][0]["request_id"] == "req" def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: @@ -78,8 +99,11 @@ def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: talker._codec_eos_id = 7 talker._force_eos_rows = None talker._mask_eos_rows = None + talker._pending_force_eos_rows = None + talker._penalty_histories = None talker._request_audio_states = {} talker._deferred_cleanup_ids = set() + talker._tts_config = ConditionalChatTTSConfig() talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)]) with torch.no_grad(): talker.head_code[0].weight.copy_(torch.eye(8, 2)) @@ -239,7 +263,9 @@ def test_talker_sample_blanks_penalty_prompt_without_touching_runner_state(mocke class _Recorder: def __call__(self, logits, sampling_metadata): - captured["prompt"] = sampling_metadata.prompt_token_ids + prompt = sampling_metadata.prompt_token_ids + assert prompt is not None + captured["prompt"] = prompt return None mocker.patch( @@ -249,6 +275,7 @@ def __call__(self, logits, sampling_metadata): talker.sample(torch.zeros(1, 8), metadata) assert captured["prompt"].tolist() == [[8, 8, 8]] + assert metadata.prompt_token_ids is not None assert metadata.prompt_token_ids.tolist() == [[0, 2, 0]] @@ -278,6 +305,150 @@ def test_compute_logits_forces_eos_for_empty_speech() -> None: assert int(logits.argmax(dim=-1)) == 7 +def _capture_sampler_logits(mocker) -> dict[str, Any]: + captured: dict[str, Any] = {} + + def _record(logits, sampling_metadata): + captured["logits"] = logits.clone() + captured["metadata"] = sampling_metadata + return _SamplerOutput(sampled_token_ids=logits.argmax(dim=-1, keepdim=True)) + + mocker.patch( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler", + return_value=_record, + ) + return captured + + +def test_sample_scores_the_upstream_windowed_codec_penalty(mocker) -> None: + talker = _make_talker() + talker._request_audio_states["req-pen"] = {"finished": False, "step": 0} + # ``head_code`` is eye(8, 2) here, so only codes 0 and 1 carry a non-zero + # logit and are the only ones a penalty can visibly move. + hidden = torch.tensor([[1.0, -2.0]]) + raw = talker.head_code[0](hidden).float() + + output = talker.make_omni_output( + hidden, + model_intermediate_buffer=[{"request_id": "req-pen", "codes": {"audio": torch.tensor([[1]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + captured = _capture_sampler_logits(mocker) + talker.sample(logits, _CodecSamplingMetadata(repetition_penalties=torch.tensor([1.05]))) + + expected = _reference_repetition_penalty(raw, torch.tensor([1]), penalty=1.05, window_size=_CODEC_PENALTY_WINDOW) + torch.testing.assert_close(captured["logits"], expected) + # A negative logit is pushed further down, so the row really did change. + assert captured["logits"][0, 1].item() < raw[0, 1].item() + # Scored once: the Sampler's own whole-stream pass is neutralized. + assert captured["metadata"].repetition_penalties.tolist() == [1.0] + # Consumed once, so a later step cannot rescore a stale row. + assert talker._penalty_histories is None + + +def test_codec_penalty_forgets_codes_older_than_the_upstream_window(mocker) -> None: + # vLLM's sampler penalty spans the whole stream, so every code ever sampled + # stays taxed and a long tail drifts toward never-sampled codes. Upstream + # scores only the last ``past_window`` frames. + talker = _make_talker() + talker._request_audio_states["req-window"] = {"finished": False, "step": 0} + hidden = torch.tensor([[1.0, -2.0]]) + raw = talker.head_code[0](hidden).float() + captured = _capture_sampler_logits(mocker) + + def emit(code: int) -> torch.Tensor: + output = talker.make_omni_output( + hidden, + model_intermediate_buffer=[{"request_id": "req-window", "codes": {"audio": torch.tensor([[code]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + talker.sample(logits, _CodecSamplingMetadata(repetition_penalties=torch.tensor([1.05]))) + return captured["logits"] + + assert emit(1)[0, 1].item() < raw[0, 1].item() + # Push code 1 out of the window with a full window of code 0. + for _ in range(_CODEC_PENALTY_WINDOW): + scored = emit(0) + + assert talker._request_audio_states["req-window"]["recent_codes"] == [0] * _CODEC_PENALTY_WINDOW + torch.testing.assert_close(scored[0, 1], raw[0, 1]) + + +def test_sample_honours_a_request_that_disabled_penalties(mocker) -> None: + talker = _make_talker() + talker._request_audio_states["req-none"] = {"finished": False, "step": 0} + hidden = torch.tensor([[1.0, -2.0]]) + raw = talker.head_code[0](hidden).float() + + output = talker.make_omni_output( + hidden, + model_intermediate_buffer=[{"request_id": "req-none", "codes": {"audio": torch.tensor([[1]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + captured = _capture_sampler_logits(mocker) + talker.sample(logits, _CodecSamplingMetadata(repetition_penalties=torch.tensor([1.0]), no_penalties=True)) + + torch.testing.assert_close(captured["logits"], raw) + + +def test_sample_restores_forced_eos_that_min_tokens_masked_away(mocker) -> None: + # ``min_tokens`` masks every stage stop_token_id, including the codec EOS + # compute_logits just forced, so the sampler returns an arbitrary id for a + # request the model already finished. + talker = _make_talker() + talker._request_audio_states["req-empty"] = {"finished": True} + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-empty", "codes": {"audio": torch.empty(0)}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + logits[0, talker._codec_eos_id] = float("-inf") + + class _MinTokensSampler: + def __call__(self, logits, sampling_metadata): + del sampling_metadata + return _SamplerOutput(sampled_token_ids=logits.argmax(dim=-1, keepdim=True)) + + mocker.patch( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler", + return_value=_MinTokensSampler(), + ) + sampled = talker.sample(logits, _PenaltyMetadata(prompt_token_ids=None)) + + assert sampled.sampled_token_ids.tolist() == [[talker._codec_eos_id]] + + +def test_sample_leaves_unforced_rows_to_the_sampler(mocker) -> None: + talker = _make_talker() + talker._request_audio_states["req-a"] = {"finished": True} + talker._request_audio_states["req-b"] = {"finished": False, "step": 0} + + output = talker.make_omni_output( + torch.ones(2, 2), + model_intermediate_buffer=[ + {"request_id": "req-a", "codes": {"audio": torch.empty(0)}}, + {"request_id": "req-b", "codes": {"audio": torch.tensor([[3]])}}, + ], + request_token_spans=[(0, 1), (1, 2)], + ) + talker.compute_logits(output.text_hidden_states) + + mocker.patch( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler", + return_value=lambda logits, sampling_metadata: _SamplerOutput(sampled_token_ids=torch.tensor([[0], [4]])), + ) + sampled = talker.sample(torch.zeros(2, 8), _PenaltyMetadata(prompt_token_ids=None)) + + assert sampled.sampled_token_ids.tolist() == [[talker._codec_eos_id], [4]] + # A single compute_logits decision must not leak into the next sample(). + assert talker._pending_force_eos_rows is None + + def test_native_duplex_chunk_budget_masks_eos_until_generate_chunk_is_full() -> None: assert _native_duplex_chunk_budget({}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, _DUPLEX_CODEC_TOKENS_PER_CHUNK) assert _native_duplex_chunk_budget({"turn_start": True}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, 0) @@ -325,21 +496,116 @@ def test_native_duplex_chunk_budget_forces_eos_after_twenty_five_frames() -> Non assert output.multimodal_outputs["meta"]["finished"][0].item() is True -def test_simplex_talker_does_not_force_eos_after_a_duplex_sized_chunk() -> None: +def test_offline_prefill_caps_codec_budget_to_the_remaining_talker_context() -> None: talker = _make_talker() - talker._request_audio_states["req-simplex"] = {"finished": False, "step": 0} + talker.emb_text = nn.Embedding(1, 2) + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + talker._text_eos_id = 0 + talker._tts_bos_id = 0 + # Long enough that the context, not upstream's max_new_token, is the binding + # limit; the Talker must never plan past its own position embeddings. + prompt_len = int(talker._tts_config.max_position_embeddings) - _OFFLINE_CODEC_MAX_NEW_TOKENS + 100 + + _, _, updates = talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_prompt_len=prompt_len, + request_id="req-offline-cap", + tts_token_ids=torch.empty(0, dtype=torch.long), + tts_hidden_states=torch.empty(0, 2), + ) + + remaining = int(talker._tts_config.max_position_embeddings) - prompt_len + assert remaining < _OFFLINE_CODEC_MAX_NEW_TOKENS + assert updates["audio_state"]["max_tokens"] == remaining + assert updates["audio_state"]["min_tokens"] is None + + +def test_talker_masks_eos_until_recorded_min_tokens() -> None: + talker = _make_talker() + talker._request_audio_states["req-early"] = { + "finished": False, + "step": 0, + "max_tokens": 8, + "min_tokens": 4, + } + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-early", "codes": {"audio": torch.tensor([[3]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert talker._request_audio_states["req-early"]["step"] == 1 + assert logits[0, 7].item() == float("-inf") + assert torch.isfinite(logits[0, 0]) + assert output.multimodal_outputs["meta"]["finished"][0].item() is False + + +def test_offline_talker_forces_eos_at_max_new_token() -> None: + talker = _make_talker() + max_tokens = 8 + talker._request_audio_states["req-cap"] = { + "finished": False, + "step": max_tokens - 2, + "max_tokens": max_tokens, + } + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-cap", "codes": {"audio": torch.tensor([[3]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert talker._request_audio_states["req-cap"]["step"] == max_tokens - 1 + assert talker._request_audio_states["req-cap"]["finished"] is True + assert int(logits.argmax(dim=-1)) == 7 + assert output.multimodal_outputs["meta"]["finished"][0].item() is True + + +def test_talker_without_a_recorded_budget_neither_forces_nor_masks_eos() -> None: + # A state rebuilt from a transported ``audio_state`` can lack budget keys; + # such rows must fall through to the raw codec head. + talker = _make_talker() + talker._request_audio_states["req-nobudget"] = {"finished": False, "step": 0} output = talker.make_omni_output( torch.tensor([[1.0, 0.0]]), - model_intermediate_buffer=[{"request_id": "req-simplex", "codes": {"audio": torch.tensor([[3]])}}], + model_intermediate_buffer=[{"request_id": "req-nobudget", "codes": {"audio": torch.tensor([[3]])}}], request_token_spans=[(0, 1)], ) logits = talker.compute_logits(output.text_hidden_states) - assert talker._request_audio_states["req-simplex"]["step"] == 0 + assert talker._request_audio_states["req-nobudget"]["step"] == 1 assert torch.allclose(logits[0], talker.head_code[0](output.text_hidden_states[0]).float()) +def test_offline_prefill_records_upstream_max_new_token_budget() -> None: + talker = _make_talker() + talker.emb_text = nn.Embedding(1, 2) + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + talker._text_eos_id = 0 + talker._tts_bos_id = 0 + + _, _, updates = talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_is_prefill=True, + request_id="req-offline", + tts_token_ids=torch.empty(0, dtype=torch.long), + tts_hidden_states=torch.empty(0, 2), + ) + + # A 2-token boundary prompt leaves far more context than upstream's + # max_new_token, so MiniCPMTTS.generate's own ceiling is what binds. + assert int(talker._tts_config.max_position_embeddings) - 2 > _OFFLINE_CODEC_MAX_NEW_TOKENS + assert updates["audio_state"]["max_tokens"] == _OFFLINE_CODEC_MAX_NEW_TOKENS + assert updates["audio_state"]["min_tokens"] is None + + def test_native_duplex_prefill_records_generate_chunk_budget() -> None: talker = _make_talker() talker.emb_text = nn.Embedding(1, 2) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index cd7705f52f7..973c4291ec8 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -66,11 +66,16 @@ stages: to_stage_2: connector_of_shared_memory # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler # (head_code logits). EOS is tts_config.eos_token_id (num_audio_tokens-1). + # Values track upstream utils.TTSSamplingParams (min_p/win_size/tau_r are + # declared there but never read, so they stay out). repetition_penalty is + # consumed by the Talker, which rescores it over upstream's 16-frame window + # instead of the Sampler's whole-stream presence penalty. + # min_tokens is upstream's min_new_token. default_sampling_params: temperature: 0.8 - top_p: 0.8 - top_k: 100 - repetition_penalty: 1.02 + top_p: 0.85 + top_k: 25 + repetition_penalty: 1.05 min_tokens: 50 max_tokens: 4096 seed: 42 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index fd56b376f96..d40715dbe44 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -54,9 +54,10 @@ stages: to_stage_2: connector_of_shared_memory default_sampling_params: temperature: 0.8 - top_p: 0.8 - top_k: 100 - repetition_penalty: 1.02 + top_p: 0.85 + top_k: 25 + # See minicpmo_4_5.yaml: the Talker rescores this over a 16-frame window. + repetition_penalty: 1.05 min_tokens: 50 max_tokens: 4096 seed: 42 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index 0bb21b1775e..26c2aeaa807 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -59,9 +59,10 @@ stages: to_stage_2: connector_of_shared_memory default_sampling_params: temperature: 0.8 - top_p: 0.8 - top_k: 100 - repetition_penalty: 1.02 + top_p: 0.85 + top_k: 25 + # See minicpmo_4_5.yaml: the Talker rescores this over a 16-frame window. + repetition_penalty: 1.05 min_tokens: 50 max_tokens: 4096 seed: 42 diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index 0734e2d96df..b7a7bc3f51d 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -71,9 +71,10 @@ stages: to_stage_2: connector_of_shared_memory default_sampling_params: temperature: 0.8 - top_p: 0.8 - top_k: 100 - repetition_penalty: 1.02 + top_p: 0.85 + top_k: 25 + # See minicpmo_4_5.yaml: the Talker rescores this over a 16-frame window. + repetition_penalty: 1.05 min_tokens: 50 max_tokens: 4096 seed: 42 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 0480176b784..68c794478cf 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from: # https://huggingface.co/openbmb/MiniCPM-o-4_5/blob/main/modeling_minicpmo.py """MiniCPM-o 4.5 native autoregressive Talker. @@ -34,6 +34,14 @@ logger = init_logger(__name__) _REPETITION_PENALTY_CHUNK_SIZE = 16 +# ``past_window`` of MiniCPMTTS's codec repetition penalty: both generate() and +# generate_chunk() build it through gen_logits(), which hardcodes +# CustomRepetitionPenaltyLogitsProcessorRepeat(penalty, num_code, 16). +_CODEC_PENALTY_WINDOW = 16 +# MiniCPMTTS.generate's max_new_token. The Talker context bounds this further; +# without it a request that never samples codec EOS keeps emitting frames for +# twice as long as upstream would, which is audible as a long silent tail. +_OFFLINE_CODEC_MAX_NEW_TOKENS = 2048 # Native duplex Talker must finish after one MiniCPMTTS.generate_chunk: # 25 codec frames (``codec_chunk_frames``) plus the terminating sample. # Without this, the single-vocab Sampler keeps the stage-1 request alive @@ -72,23 +80,31 @@ def _apply_batched_repetition_penalty( logits: torch.Tensor, histories: Sequence[torch.Tensor], *, - penalty: float, + penalty: float | torch.Tensor, window_size: int, ) -> torch.Tensor: - """Apply request-local frequency penalties to a batch of codec logits.""" + """Apply request-local frequency penalties to a batch of codec logits. + + ``penalty`` may be a scalar or one value per row, mirroring upstream's + per-request ``sampling_params.repetition_penalty``. + """ if logits.ndim != 2: raise ValueError(f"batched codec logits must be 2D, got shape {tuple(logits.shape)}") batch_size, vocab_size = logits.shape if len(histories) != batch_size: raise ValueError(f"expected {batch_size} codec histories, got {len(histories)}") - if penalty == 1.0: + if batch_size == 0: return logits - if batch_size == 0: + penalties = torch.as_tensor(penalty, device=logits.device, dtype=logits.dtype).reshape(-1) + if penalties.numel() == 1: + penalties = penalties.expand(batch_size) + elif penalties.numel() != batch_size: + raise ValueError(f"expected 1 or {batch_size} codec repetition penalties, got {penalties.numel()}") + if not bool((penalties != 1.0).any()): return logits penalized = logits.clone() - penalty_tensor = torch.as_tensor(penalty, device=logits.device, dtype=logits.dtype) for start in range(0, batch_size, _REPETITION_PENALTY_CHUNK_SIZE): end = min(start + _REPETITION_PENALTY_CHUNK_SIZE, batch_size) chunk_logits = logits[start:end] @@ -106,7 +122,7 @@ def _apply_batched_repetition_penalty( encoded, minlength=(end - start) * vocab_size, ).reshape(end - start, vocab_size) - alpha = torch.pow(penalty_tensor, frequencies.to(dtype=logits.dtype)) + alpha = torch.pow(penalties[start:end].unsqueeze(1), frequencies.to(dtype=logits.dtype)) penalized[start:end] = torch.where(chunk_logits < 0, chunk_logits * alpha, chunk_logits / alpha) return penalized @@ -139,6 +155,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.vllm_config = vllm_config self._force_eos_rows: list[bool] | None = None self._mask_eos_rows: list[bool] | None = None + self._pending_force_eos_rows: list[bool] | None = None + self._penalty_histories: list[torch.Tensor] | None = None self._request_audio_states: dict[str, dict[str, Any]] = {} self._deferred_cleanup_ids: set[str] = set() @@ -315,11 +333,21 @@ def preprocess( f"tts_ids={token_ids.shape[0]} tts_hidden={hidden_states.shape[0]} " f"prompt_len={info_dict.get('_omni_prompt_len')}" ) - state: dict[str, Any] = {"finished": empty_condition, "step": 0} if native_duplex: max_tokens, min_tokens = _native_duplex_chunk_budget(meta if isinstance(meta, Mapping) else None) - state["max_tokens"] = max_tokens - state["min_tokens"] = min_tokens + else: + # MiniCPMTTS.generate()'s max_new_token, clamped to what the + # Talker context can still hold. Sampler min_tokens (upstream's + # min_new_token=50) comes from the deploy YAML. + remaining = int(self._tts_config.max_position_embeddings) - target_len + max_tokens = max(min(_OFFLINE_CODEC_MAX_NEW_TOKENS, remaining), 1) + min_tokens = None + state: dict[str, Any] = { + "finished": empty_condition, + "step": 0, + "max_tokens": max_tokens, + "min_tokens": min_tokens, + } request_states = getattr(self, "_request_audio_states", None) if request_states is None: request_states = {} @@ -354,7 +382,8 @@ def preprocess( # hand the same id to make_omni_output so Code2Wav sees it this step. code = input_ids.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1)[-1:] embeds = self.emb_code[0](code) - if int(code.item()) == int(self._codec_eos_id): + code_id = int(code.item()) + if code_id == int(self._codec_eos_id): if isinstance(state, dict): state["finished"] = True elif stored is None: @@ -388,6 +417,8 @@ def make_omni_output( terminal_flags = [torch.tensor(False, dtype=torch.bool) for _ in infos] force_eos_rows = [False] * len(infos) mask_eos_rows = [False] * len(infos) + empty_history = hidden.new_empty((0,), dtype=torch.long) + penalty_histories = [empty_history for _ in infos] for index, info in enumerate(infos): info_dict = info if isinstance(info, dict) else {} native_duplex = info_dict.get("native_duplex") is True @@ -444,13 +475,22 @@ def make_omni_output( audio = codes.get("audio") if isinstance(codes, Mapping) else None if isinstance(audio, torch.Tensor) and audio.numel() > 0: codec_deltas[index] = audio.to(device=hidden.device, dtype=torch.long).reshape(-1, 1) - if state.get("max_tokens") is not None: - state["step"] = int(state.get("step", 0)) + 1 + state["step"] = int(state.get("step", 0)) + 1 + # ``audio`` is the id sampled last step, i.e. exactly upstream's + # ``new_tokens[:, 0:t]`` history for the logits computed below. + recent = state.get("recent_codes") + recent = (recent if isinstance(recent, list) else []) + codec_deltas[index].reshape(-1).tolist() + state["recent_codes"] = recent[-_CODEC_PENALTY_WINDOW:] + recent_codes = state.get("recent_codes") + if recent_codes: + penalty_histories[index] = torch.tensor(recent_codes, dtype=torch.long, device=hidden.device) max_tokens = state.get("max_tokens") min_tokens = state.get("min_tokens") step = int(state.get("step", 0)) - # 26 samples include the terminating EOS, so force it after 25 - # forwarded frames — the same cadence as generate_chunk. + # Duplex: 26 samples include the terminating EOS, so force it after + # 25 forwarded frames — the same cadence as generate_chunk. + # Offline: force-stop at the remaining Talker context rather than + # waiting for a sampled EOS. hit_chunk_limit = max_tokens is not None and step >= int(max_tokens) - 1 chunk_done = empty_speech or hit_chunk_limit if hit_chunk_limit: @@ -462,11 +502,13 @@ def make_omni_output( mask_eos_rows[index] = not force_eos_rows[index] and min_tokens is not None and step < int(min_tokens) terminal_flags[index] = torch.tensor(chunk_done, dtype=torch.bool) - # Empty-speech and finished duplex chunks must sample codec EOS so - # the scheduler releases the request; mid-chunk rows mask EOS until - # the 26-token generate_chunk budget is met. + # Empty-speech rows, finished duplex chunks, and offline requests that + # fill the remaining Talker context must sample codec EOS so the + # scheduler releases the request. Mid-chunk rows mask EOS until + # min_tokens. self._force_eos_rows = force_eos_rows self._mask_eos_rows = mask_eos_rows + self._penalty_histories = penalty_histories meta_outputs = {"finished": terminal_flags} if emit_duplex_metadata: meta_outputs.update( @@ -549,6 +591,11 @@ def compute_logits(self, hidden_states, *args, **kwargs): self._mask_eos_rows = None need_force = bool(force_eos and len(force_eos) == logits.shape[0] and any(force_eos)) need_mask = bool(mask_eos and len(mask_eos) == logits.shape[0] and any(mask_eos)) + # sample() re-applies the decision on the sampled ids: vLLM's + # MinTokensLogitsProcessor runs after this and would blank the codec EOS + # we just forced (it is in the stage's ``stop_token_ids``), leaving an + # all -inf row and a request that never releases. + self._pending_force_eos_rows = force_eos if need_force else None if need_force or need_mask: logits = logits.clone() eos_id = int(self._codec_eos_id) @@ -572,7 +619,64 @@ def sample(self, logits, sampling_metadata): sampling_metadata, prompt_token_ids=blank_scheduler_prompt_for_penalties(prompt_ids, logits.shape[-1]), ) - return Sampler()(logits, sampling_metadata) + logits, sampling_metadata = self._apply_codec_repetition_penalty(logits, sampling_metadata) + force_eos = self._pending_force_eos_rows + self._pending_force_eos_rows = None + output = Sampler()(logits, sampling_metadata) + return self._force_eos_on_sampled_ids(output, force_eos) + + def _apply_codec_repetition_penalty(self, logits, sampling_metadata): + """Score MiniCPMTTS.generate's windowed codec penalty, not vLLM's. + + Upstream taxes a code by ``penalty ** frequency`` over the last + ``past_window`` frames only (``gen_logits`` builds + ``CustomRepetitionPenaltyLogitsProcessorRepeat(penalty, num_code, 16)``). + vLLM's is presence-based over the whole stream, so on a codec stream + thousands of frames long every code ever sampled ends up taxed by the + same flat factor while never-sampled codes stay untouched, and the tail + of a long answer drifts off the speech manifold into near-silence. + + Runs before ``Sampler`` so the penalty lands ahead of top-k/top-p, as + upstream does. Upstream scores it after dividing by temperature, but + the penalty only rescales and preserves sign, so the two orders agree. + """ + histories = self._penalty_histories + self._penalty_histories = None + penalties = getattr(sampling_metadata, "repetition_penalties", None) + if ( + not isinstance(logits, torch.Tensor) + or histories is None + or len(histories) != logits.shape[0] + or not isinstance(penalties, torch.Tensor) + or getattr(sampling_metadata, "no_penalties", False) + ): + return logits, sampling_metadata + logits = _apply_batched_repetition_penalty( + logits, + histories, + penalty=penalties.to(device=logits.device, dtype=logits.dtype), + window_size=_CODEC_PENALTY_WINDOW, + ) + # Neutralize the sampler's own pass so the penalty is scored once. + return logits, replace(sampling_metadata, repetition_penalties=torch.ones_like(penalties)) + + def _force_eos_on_sampled_ids(self, output: Any, force_eos: list[bool] | None) -> Any: + """Overwrite sampled ids for rows the model terminated this step. + + The codec EOS is a stage ``stop_token_ids`` entry, so vLLM's + ``min_tokens`` processor masks it for the first ``min_tokens`` steps. + A row the model forced to EOS therefore reaches the sampler as all + -inf and comes back as an arbitrary codec id, which keeps an + already-finished request decoding until its length cap. + """ + if not force_eos or not any(force_eos): + return output + sampled = getattr(output, "sampled_token_ids", None) + if not isinstance(sampled, torch.Tensor) or sampled.shape[0] != len(force_eos): + return output + rows = torch.tensor(force_eos, dtype=torch.bool, device=sampled.device) + sampled[rows] = int(self._codec_eos_id) + return output def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): return self._load_native_weights(weights)