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 975f2cdaa..8a7593623 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 @@ -12,7 +12,9 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( MiniCPMO45OmniForConditionalGeneration, ) from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + _REPETITION_PENALTY_CHUNK_SIZE, MiniCPMO45OmniTTSForConditionalGeneration, + _apply_batched_repetition_penalty, _max_audio_tokens, _restore_weight_norm_weight, ) @@ -73,6 +75,50 @@ def _routed(output, index: int): ) +def _reference_repetition_penalty( + logits: torch.Tensor, + history: torch.Tensor, + *, + penalty: float, + window_size: int, +) -> torch.Tensor: + if penalty == 1.0 or history.numel() == 0: + return logits + recent = history.reshape(-1)[-window_size:].to(device=logits.device, dtype=torch.long) + frequencies = torch.bincount(recent, minlength=logits.shape[-1]).to(dtype=logits.dtype) + alpha = torch.pow(torch.as_tensor(penalty, device=logits.device, dtype=logits.dtype), frequencies) + return torch.where(logits < 0, logits * alpha, logits / alpha) + + +def _make_sampling_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: + talker = _make_talker() + talker._codec_temperature = 0.8 + talker._codec_top_k = 5 + talker._codec_top_p = 0.85 + talker._codec_repetition_penalty = 1.05 + talker._request_audio_states = { + "req-a": {"min_tokens": 2}, + "req-b": {"min_tokens": 0}, + } + talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)]) + with torch.no_grad(): + talker.head_code[0].weight.copy_( + torch.tensor( + [ + [0.1, -0.3], + [0.2, 0.4], + [-0.5, 0.2], + [0.7, -0.1], + [-0.2, 0.8], + [0.6, 0.3], + [-0.4, -0.7], + [0.3, 0.5], + ] + ) + ) + return talker + + @pytest.mark.parametrize( ("condition_tokens", "expected"), [(3, 64), (100, 1000), (1000, 2048)], @@ -101,16 +147,153 @@ def test_weight_norm_restore_matches_checkpoint_parametrization_in_bfloat16() -> assert torch.equal(restored, linear.weight) +def test_batched_repetition_penalty_matches_request_local_rows() -> None: + logits = torch.tensor( + [ + [-2.0, -1.0, 1.0, 2.0, 3.0, 4.0], + [4.0, 3.0, 2.0, 1.0, -1.0, -2.0], + [1.0, -1.0, 2.0, -2.0, 3.0, -3.0], + ] + ) + histories = [ + torch.tensor([1, 1, 2, 5]), + torch.tensor([0, 4, 4]), + torch.empty(0, dtype=torch.long), + ] + + actual = _apply_batched_repetition_penalty( + logits, + histories, + penalty=1.2, + window_size=3, + ) + expected = torch.cat( + [ + _reference_repetition_penalty( + logits[row : row + 1], + history, + penalty=1.2, + window_size=3, + ) + for row, history in enumerate(histories) + ], + dim=0, + ) + + assert torch.equal(actual, expected) + + +def test_batched_repetition_penalty_matches_rows_across_chunks(mocker) -> None: + batch_size = 2 * _REPETITION_PENALTY_CHUNK_SIZE + 1 + vocab_size = 11 + logits = torch.arange(batch_size * vocab_size, dtype=torch.float32).reshape(batch_size, vocab_size) - 100 + histories = [ + torch.tensor([(row + offset) % vocab_size for offset in range(row % 7)], dtype=torch.long) + for row in range(batch_size) + ] + + expected = torch.cat( + [ + _reference_repetition_penalty( + logits[row : row + 1], + history, + penalty=1.2, + window_size=5, + ) + for row, history in enumerate(histories) + ], + dim=0, + ) + bincount = mocker.spy(torch, "bincount") + actual = _apply_batched_repetition_penalty( + logits, + histories, + penalty=1.2, + window_size=5, + ) + + assert torch.equal(actual, expected) + assert [call.kwargs["minlength"] for call in bincount.call_args_list] == [ + _REPETITION_PENALTY_CHUNK_SIZE * vocab_size, + _REPETITION_PENALTY_CHUNK_SIZE * vocab_size, + vocab_size, + ] + + +def test_batched_codec_sampling_is_request_order_independent() -> None: + talker = _make_sampling_talker() + histories = [torch.tensor([1, 1, 2]), torch.tensor([3, 4, 4])] + hidden = torch.tensor([[0.25, -0.75], [-0.5, 0.5]]) + + forward_order = talker._sample_audio_codes( + hidden, + histories, + ["req-a", "req-b"], + [0, 0], + ) + talker._request_generators = {} + reverse_order = talker._sample_audio_codes( + hidden.flip(0), + list(reversed(histories)), + ["req-b", "req-a"], + [0, 0], + ) + + assert torch.equal(forward_order[0], reverse_order[1]) + assert torch.equal(forward_order[1], reverse_order[0]) + + +def test_batched_codec_sampling_preserves_request_rng_after_compaction() -> None: + talker = _make_sampling_talker() + history_a = torch.tensor([1, 1, 2]) + history_b = torch.tensor([3, 4, 4]) + hidden_a = torch.tensor([[0.25, -0.75]]) + hidden_b = torch.tensor([[-0.5, 0.5]]) + + first_batch = talker._sample_audio_codes( + torch.cat([hidden_a, hidden_b]), + [history_a, history_b], + ["req-a", "req-b"], + [0, 0], + ) + compacted_b = talker._sample_audio_codes( + hidden_b, + [history_b], + ["req-b"], + [1], + ) + + talker._request_generators = {} + standalone_b_first = talker._sample_audio_codes( + hidden_b, + [history_b], + ["req-b"], + [0], + )[0] + standalone_b_second = talker._sample_audio_codes( + hidden_b, + [history_b], + ["req-b"], + [1], + )[0] + + assert torch.equal(first_batch[1], standalone_b_first) + assert torch.equal(compacted_b[0], standalone_b_second) + + def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> None: talker = _make_talker() seen: list[tuple[str, list[float], list[int]]] = [] - def sample(hidden, history, request_id, step): - assert step == 0 - seen.append((request_id, hidden.reshape(-1).tolist(), history.tolist())) - return torch.tensor(2 if request_id == "req-a" else 3) + def sample(hidden, histories, request_ids, steps): + assert steps == [0, 0] + seen.extend( + (request_id, hidden[row].reshape(-1).tolist(), histories[row].tolist()) + for row, request_id in enumerate(request_ids) + ) + return torch.tensor([2, 3]) - mocker.patch.object(talker, "_sample_audio_code", side_effect=sample) + batched_sample = mocker.patch.object(talker, "_sample_audio_codes", side_effect=sample) infos = [ {"request_id": "req-a", "audio_codes": {"accumulated": torch.tensor([1])}}, {"request_id": "req-b", "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}}, @@ -126,6 +309,7 @@ def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> N ("req-a", [2.0, 0.0], [1]), ("req-b", [3.0, 0.0], []), ] + batched_sample.assert_called_once() assert infos[0]["audio_codes"]["accumulated"].tolist() == [1, 2] assert infos[1]["audio_codes"]["accumulated"].tolist() == [3] assert set(output.multimodal_outputs) == {"codes", "meta"} @@ -140,7 +324,7 @@ def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> N def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None: talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2)) + mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2, 2])) infos = [ { "request_id": "req-a", @@ -186,7 +370,7 @@ def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None: def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None: talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2)) + mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2])) info = { "request_id": "req-missing-fence", "native_duplex": True, @@ -203,7 +387,7 @@ def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None: def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> None: talker = _make_talker() - sample = mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2)) + sample = mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2])) infos = [ { "request_id": "req-prefill", @@ -225,7 +409,7 @@ def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> ) sample.assert_called_once() - assert sample.call_args.args[2] == "req-decode" + assert sample.call_args.args[2] == ["req-decode"] assert infos[0]["audio_state"]["step"] == 0 assert infos[0]["audio_codes"]["accumulated"].numel() == 0 assert infos[1]["audio_state"]["step"] == 5 @@ -235,7 +419,7 @@ def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None: talker = _make_talker() - sample = mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(7)) + sample = mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([7])) info = { "request_id": "req-stop", "audio_state": {"step": 3}, @@ -265,7 +449,7 @@ def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None: def test_max_token_terminal_drops_unconsumed_codec_delta(mocker) -> None: talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(3)) + mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([3])) info = { "request_id": "req-limit", "audio_state": {"step": 1, "max_tokens": 2}, @@ -288,7 +472,7 @@ def test_max_token_terminal_drops_unconsumed_codec_delta(mocker) -> None: def test_request_local_state_survives_missing_runner_buffer_update(mocker) -> None: talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(3)) + mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([3])) first_info = { "request_id": "req-local-state", "audio_state": {"step": 1, "max_tokens": 3}, 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 02b0bf91f..9e9dfe952 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 @@ -11,7 +11,8 @@ Pipeline: 4. Continuously generate request-aligned discrete audio-code deltas """ -from collections.abc import Iterable +from collections.abc import Iterable, Sequence +from dataclasses import dataclass from typing import Any import torch @@ -32,6 +33,7 @@ from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) _REPETITION_WINDOW = 16 +_REPETITION_PENALTY_CHUNK_SIZE = 16 _MIN_AUDIO_TOKENS = 64 _MAX_AUDIO_TOKENS = 2048 _AUDIO_TOKENS_PER_TEXT_TOKEN = 10 @@ -46,6 +48,17 @@ _CODEC_MIN_TOKENS = 50 _DUPLEX_CODEC_TOKENS_PER_CHUNK = 26 +@dataclass(slots=True) +class _PendingCodecSample: + output_index: int + hidden_row: torch.Tensor + codes: torch.Tensor + request_id: str + step: int + state: dict[str, Any] + info: dict[str, Any] + + def _max_audio_tokens(condition_tokens: int) -> int: """Bound codec generation with a conservative text-length estimate. @@ -65,20 +78,48 @@ def _restore_weight_norm_weight(weight_g: torch.Tensor, weight_v: torch.Tensor) return torch._weight_norm(weight_v, weight_g, dim=0) -def _apply_repetition_penalty( +def _apply_batched_repetition_penalty( logits: torch.Tensor, - history: torch.Tensor, + histories: Sequence[torch.Tensor], *, penalty: float, window_size: int, ) -> torch.Tensor: - """Match MiniCPMTTS' frequency-aware repetition penalty.""" - if penalty == 1.0 or history.numel() == 0: + """Apply request-local frequency penalties to a batch of codec logits.""" + 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: return logits - recent = history.reshape(-1)[-window_size:].to(device=logits.device, dtype=torch.long) - frequencies = torch.bincount(recent, minlength=logits.shape[-1]).to(dtype=logits.dtype) - alpha = torch.pow(torch.as_tensor(penalty, device=logits.device, dtype=logits.dtype), frequencies) - return torch.where(logits < 0, logits * alpha, logits / alpha) + + if batch_size == 0: + 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] + encoded_rows: list[torch.Tensor] = [] + for local_row, history in enumerate(histories[start:end]): + recent = history.reshape(-1)[-window_size:].to(device=logits.device, dtype=torch.long) + if recent.numel() > 0: + encoded_rows.append(recent + local_row * vocab_size) + if not encoded_rows: + continue + + # Bound the int64 bincount workspace independently of request concurrency. + encoded = encoded_rows[0] if len(encoded_rows) == 1 else torch.cat(encoded_rows) + frequencies = torch.bincount( + encoded, + minlength=(end - start) * vocab_size, + ).reshape(end - start, vocab_size) + alpha = torch.pow(penalty_tensor, frequencies.to(dtype=logits.dtype)) + penalized[start:end] = torch.where(chunk_logits < 0, chunk_logits * alpha, chunk_logits / alpha) + + return penalized def _apply_top_k_top_p( @@ -370,28 +411,47 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): self._request_generators[request_id] = generator return generator - def _sample_audio_code( + def _sample_audio_codes( self, - hidden_state: torch.Tensor, - history: torch.Tensor, - request_id: str, - step: int, + hidden_states: torch.Tensor, + histories: Sequence[torch.Tensor], + request_ids: Sequence[str], + steps: Sequence[int], ) -> torch.Tensor: - logits = self.head_code[0](hidden_state).float() / self._codec_temperature + batch_size = int(hidden_states.shape[0]) + if not (len(histories) == len(request_ids) == len(steps) == batch_size): + raise ValueError( + "MiniCPM-o batched codec sampling requires one history, request id, " + f"and step per hidden row, got batch={batch_size}, histories={len(histories)}, " + f"request_ids={len(request_ids)}, steps={len(steps)}" + ) + if batch_size == 0: + return torch.empty(0, dtype=torch.long, device=hidden_states.device) + + logits = self.head_code[0](hidden_states).float() / self._codec_temperature eos_id = self._num_audio_tokens - 1 - logits = _apply_repetition_penalty( + logits = _apply_batched_repetition_penalty( logits, - history, + histories, penalty=self._codec_repetition_penalty, window_size=_REPETITION_WINDOW, ) request_states = getattr(self, "_request_audio_states", {}) - state = request_states.get(request_id) - min_tokens = ( - int(state.get("min_tokens", self._codec_min_tokens)) if isinstance(state, dict) else self._codec_min_tokens + mask_eos_values: list[bool] = [] + for request_id, step in zip(request_ids, steps, strict=True): + state = request_states.get(request_id) + min_tokens = ( + int(state.get("min_tokens", self._codec_min_tokens)) + if isinstance(state, dict) + else self._codec_min_tokens + ) + mask_eos_values.append(step < min_tokens) + mask_eos = torch.tensor( + mask_eos_values, + dtype=torch.bool, + device=logits.device, ) - if step < min_tokens: - logits[..., eos_id] = float("-inf") + logits[:, eos_id].masked_fill_(mask_eos, float("-inf")) logits = _apply_top_k_top_p( logits, top_k=self._codec_top_k, @@ -399,11 +459,21 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): min_tokens_to_keep=3, ) probabilities = torch.softmax(logits, dim=-1) - return torch.multinomial( - probabilities, - num_samples=1, - generator=self._request_generator(request_id, probabilities.device), - ).reshape(()) + # torch.multinomial accepts one generator for the entire call. Keep it + # per row so compaction, request ordering, or another request finishing + # cannot perturb request-local RNG streams; the expensive projection + # and logit transforms above remain fully batched. + return torch.cat( + [ + torch.multinomial( + probabilities[row : row + 1], + num_samples=1, + generator=self._request_generator(request_id, probabilities.device), + ) + for row, request_id in enumerate(request_ids) + ], + dim=0, + ).reshape(-1) def make_omni_output( self, @@ -426,15 +496,18 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): ) emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos) - stop_rows: list[torch.Tensor] = [] - codec_deltas: list[torch.Tensor] = [] - terminal_flags: list[torch.Tensor] = [] + # Rows default to continue. Only previously finished or newly terminal + # requests stop; prefill/ineligible rows stay aligned as False. + stop_flags = [False] * len(infos) native_duplex_flags: list[torch.Tensor] = [] duplex_epochs: list[torch.Tensor] = [] duplex_turn_ids: list[torch.Tensor] = [] segment_texts_utf8: list[torch.Tensor] = [] turn_end_flags: list[torch.Tensor] = [] empty_delta = hidden.new_empty((0, 1), dtype=torch.long) + codec_deltas = [empty_delta for _ in infos] + terminal_flags = [torch.tensor(False, dtype=torch.bool) for _ in infos] + pending_samples: list[_PendingCodecSample] = [] for index, info in enumerate(infos): info_dict = info if isinstance(info, dict) else {} native_duplex = info_dict.get("native_duplex") is True @@ -480,16 +553,10 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): turn_end_flags.append(torch.tensor(native_duplex and contains_turn_eos, dtype=torch.bool)) if not isinstance(info, dict): - stop_rows.append(hidden.new_tensor([0.0, float("-inf")])) - codec_deltas.append(empty_delta) - terminal_flags.append(torch.tensor(False, dtype=torch.bool)) continue start, end = spans[index] end = min(int(end), int(hidden.shape[0])) if int(start) >= end: - stop_rows.append(hidden.new_tensor([0.0, float("-inf")])) - codec_deltas.append(empty_delta) - terminal_flags.append(torch.tensor(False, dtype=torch.bool)) continue request_id = str(info.get("request_id", index)) request_states = getattr(self, "_request_audio_states", None) @@ -501,17 +568,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): state = dict(info.get("audio_state", {}) or {}) request_states[request_id] = state if state.get("finished"): - stop_rows.append(hidden.new_tensor([float("-inf"), 0.0])) - codec_deltas.append(empty_delta) - terminal_flags.append(torch.tensor(False, dtype=torch.bool)) + stop_flags[index] = True continue if not sample_eligible[index]: # vLLM computes a logit row for incomplete chunked prefills but # discards its sampled token. Advancing codec/RNG state here # would make output depend on prefill chunking and compaction. - stop_rows.append(hidden.new_tensor([0.0, float("-inf")])) - codec_deltas.append(empty_delta) - terminal_flags.append(torch.tensor(False, dtype=torch.bool)) continue codes = state.get("codes") if not isinstance(codes, torch.Tensor): @@ -521,8 +583,42 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): else: codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1) step = int(state.get("step", 0)) - sampled = self._sample_audio_code(hidden[end - 1 : end], codes, request_id, step) - sampled_id = int(sampled.item()) + pending_samples.append( + _PendingCodecSample( + output_index=index, + hidden_row=hidden[end - 1 : end], + codes=codes, + request_id=request_id, + step=step, + state=state, + info=info, + ) + ) + + if pending_samples: + active_hidden_rows = [pending.hidden_row for pending in pending_samples] + active_hidden = ( + active_hidden_rows[0] if len(active_hidden_rows) == 1 else torch.cat(active_hidden_rows, dim=0) + ) + sampled_batch = self._sample_audio_codes( + active_hidden, + [pending.codes for pending in pending_samples], + [pending.request_id for pending in pending_samples], + [pending.step for pending in pending_samples], + ) + # One batched device-to-host synchronization replaces one .item() + # synchronization per request. + sampled_ids = sampled_batch.detach().to(device="cpu").tolist() + else: + sampled_batch = hidden.new_empty((0,), dtype=torch.long) + sampled_ids = [] + + for row, pending in enumerate(pending_samples): + sampled = sampled_batch[row].reshape(()) + sampled_id = int(sampled_ids[row]) + codes = pending.codes + state = pending.state + info = pending.info is_eos = sampled_id == self._num_audio_tokens - 1 state["step"] = int(state.get("step", 0)) + 1 reached_limit = int(state["step"]) >= int(state.get("max_tokens", 2048)) @@ -541,11 +637,15 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): "current": sampled.reshape(1), "accumulated": codes, } - codec_deltas.append(delta) - terminal_flags.append(torch.tensor(finished, dtype=torch.bool)) - stop_rows.append(hidden.new_tensor([float("-inf"), 0.0] if finished else [0.0, float("-inf")])) + codec_deltas[pending.output_index] = delta + terminal_flags[pending.output_index] = torch.tensor(finished, dtype=torch.bool) + stop_flags[pending.output_index] = finished - self._batch_stop_logits = torch.stack(stop_rows, dim=0) if stop_rows else hidden.new_empty((0, 2)) + self._batch_stop_logits = ( + hidden.new_tensor([[float("-inf"), 0.0] if finished else [0.0, float("-inf")] for finished in stop_flags]) + if stop_flags + else hidden.new_empty((0, 2)) + ) # Lists are deliberate: the runner routes element i to request i, # preserving compaction alignment while emitting only this step's code. meta_outputs = {"finished": terminal_flags}