From 43ef1e61fbb880a84cfa3416afd56e2be06a413e Mon Sep 17 00:00:00 2001 From: jingchengtian <156733679+jingchengtian@users.noreply.github.com> Date: Wed, 30 Sep 2026 01:13:50 +0800 Subject: [PATCH] [Perf][TTS] Add per-row generator support to MOSS-TTS talker to eliminate serial fallback for seeded batched requests (#7922) Signed-off-by: jingchengtian Signed-off-by: jingchengtian Co-authored-by: jingchengtian --- .../moss_tts/test_per_row_generators.py | 120 ++++++++++++++++++ .../moss_tts/modeling_moss_tts_local.py | 46 ++++++- .../moss_tts/modeling_moss_tts_local_depth.py | 7 + .../moss_tts/modeling_moss_tts_talker.py | 8 ++ 4 files changed, 178 insertions(+), 3 deletions(-) create mode 100644 tests/model_executor/models/moss_tts/test_per_row_generators.py diff --git a/tests/model_executor/models/moss_tts/test_per_row_generators.py b/tests/model_executor/models/moss_tts/test_per_row_generators.py new file mode 100644 index 000000000..4719ec8cb --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_per_row_generators.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU regression tests for MOSS-TTS per-row talker generators.""" + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts import modeling_moss_tts_local + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts] + +BATCH = 4 +VOCAB = 1024 + + +def _logits() -> torch.Tensor: + return torch.randn(BATCH, VOCAB, generator=torch.Generator().manual_seed(7)) + + +def _row_generators() -> list[torch.Generator]: + return [torch.Generator().manual_seed(100 + row) for row in range(BATCH)] + + +def _sample( + logits: torch.Tensor, + generators: list[torch.Generator | None] | None, +) -> torch.Tensor: + return modeling_moss_tts_local._sample_token( + logits, + temperature=1.7, + top_k=25, + top_p=0.8, + do_sample=True, + generators=generators, + ) + + +def test_normalize_generators_rejects_length_mismatch() -> None: + assert modeling_moss_tts_local._normalize_generators(None, BATCH) is None + + generators = _row_generators() + normalized = modeling_moss_tts_local._normalize_generators(generators, BATCH) + assert normalized is not None + assert all(row is expected for row, expected in zip(normalized, generators)) + + with pytest.raises(ValueError, match=f"Expected {BATCH} per-row generators, but got {BATCH - 1}"): + modeling_moss_tts_local._normalize_generators(_row_generators()[:-1], BATCH) + with pytest.raises(ValueError, match=f"Expected {BATCH} per-row generators, but got {BATCH + 1}"): + modeling_moss_tts_local._normalize_generators(_row_generators() + [torch.Generator()], BATCH) + + +def test_sample_token_validates_generator_count() -> None: + """A short generators list must fail loudly instead of sampling the tail rows + from the global RNG.""" + logits = _logits() + + with pytest.raises(ValueError, match=f"Expected {BATCH} per-row generators, but got {BATCH - 1}"): + _sample(logits, _row_generators()[:-1]) + + +def test_sample_token_per_row_generators_are_reproducible() -> None: + logits = _logits() + + torch.testing.assert_close(_sample(logits, _row_generators()), _sample(logits, _row_generators())) + + +def test_sample_token_rows_do_not_depend_on_batch_composition() -> None: + """Each row must be sampled by its own generator alone, so batching a row + with other requests cannot change its output.""" + logits = _logits() + + batched = _sample(logits, _row_generators()) + for row in range(BATCH): + single = modeling_moss_tts_local._sample_token( + logits[row : row + 1], + temperature=1.7, + top_k=25, + top_p=0.8, + do_sample=True, + generators=[torch.Generator().manual_seed(100 + row)], + ) + torch.testing.assert_close(batched[row], single[0]) + + +def test_sample_token_all_none_generators_keep_batched_path() -> None: + """The runner only passes ``generators`` when a row is explicitly seeded, so + an all-``None`` list must behave exactly like the scalar batched path.""" + logits = _logits() + + with_generators = modeling_moss_tts_local._sample_token( + logits, + temperature=1.7, + top_k=25, + top_p=0.8, + do_sample=True, + generator=torch.Generator().manual_seed(42), + generators=[None] * BATCH, + ) + without_generators = modeling_moss_tts_local._sample_token( + logits, + temperature=1.7, + top_k=25, + top_p=0.8, + do_sample=True, + generator=torch.Generator().manual_seed(42), + ) + torch.testing.assert_close(with_generators, without_generators) + + +def test_sample_token_single_seeded_row_is_reproducible() -> None: + """One seeded row routes the batch through the per-row loop; that row must + stay reproducible while unseeded rows keep using the global RNG.""" + logits = _logits() + + def run() -> torch.Tensor: + generators: list[torch.Generator | None] = [None] * BATCH + generators[1] = torch.Generator().manual_seed(101) + return _sample(logits, generators) + + torch.testing.assert_close(run()[1], run()[1]) diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py index c9ce11a93..f1d23a7de 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). @@ -23,7 +26,7 @@ previous KV-cache loop -- causal attention, only the last position is read. from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence import torch import torch.nn as nn @@ -138,6 +141,24 @@ class MossTTSRealtimeLocalTransformer(nn.Module): return codes +def _normalize_generators( + generators: Sequence[torch.Generator | None] | None, + batch_size: int, +) -> list[torch.Generator | None] | None: + """Validate per-row generators against the batch they will sample for. + + Mirrors ``Qwen3CodePredictor._normalize_generators``: a length mismatch is + an error instead of a silent per-row fallback to the global RNG, which would + quietly break reproducibility of seeded requests. + """ + if generators is None: + return None + row_generators = list(generators) + if len(row_generators) != batch_size: + raise ValueError(f"Expected {batch_size} per-row generators, but got {len(row_generators)}.") + return row_generators + + def _sample_token( logits: torch.Tensor, temperature: float, @@ -145,6 +166,7 @@ def _sample_token( top_p: float, do_sample: bool, generator: torch.Generator | None = None, + generators: Sequence[torch.Generator | None] | None = None, ) -> torch.Tensor: """Top-k + top-p sampling for the upstream inference branch. @@ -153,6 +175,12 @@ def _sample_token( vocabulary. This preserves the categorical distribution when the top-k boundary has no ties, but it is not seed/bit equivalent to multinomial over a full-vocabulary tensor because random-number mapping depends on width. + + When ``generators`` is provided (per-row), each row is sampled with its + own generator via separate ``multinomial`` calls. This keeps the rest of + the forward batched while making seeded requests reproducible per-row. + ``generators`` must hold exactly one entry per batch row; a mismatch raises + ``ValueError`` rather than silently falling back to the global RNG. """ if not do_sample or temperature <= 0: return logits.argmax(dim=-1) @@ -181,8 +209,20 @@ def _sample_token( logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_indices, logits) probs = F.softmax(logits, dim=-1) - flat = probs.reshape(-1, probs.shape[-1]) - sampled = torch.multinomial(flat, num_samples=1, generator=generator).reshape(probs.shape[:-1]) + B = int(probs.shape[0]) + row_generators = _normalize_generators(generators, B) + if row_generators is not None and any(gen is not None for gen in row_generators): + rows = [] + for row in range(B): + row_probs = probs[row : row + 1] + flat = row_probs.reshape(-1, row_probs.shape[-1]) + sampled = torch.multinomial(flat, num_samples=1, generator=row_generators[row]) + rows.append(sampled) + sampled = torch.cat(rows, dim=0).reshape(probs.shape[:-1]) + else: + flat = probs.reshape(-1, probs.shape[-1]) + sampled = torch.multinomial(flat, num_samples=1, generator=generator).reshape(probs.shape[:-1]) + if compact_indices is not None: return compact_indices.gather(-1, sampled.unsqueeze(-1)).squeeze(-1) return sampled diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py index a7e7606a3..c3f226fd1 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py @@ -21,6 +21,8 @@ checkpoint 1:1 so ``load_weights()`` needs no remapping. from __future__ import annotations +from collections.abc import Sequence + import torch import torch.nn as nn import torch.nn.functional as F @@ -236,6 +238,7 @@ class MossTTSLocalDepthTransformer(nn.Module): top_p: float, do_sample: bool, generator: torch.Generator | None, + generators: Sequence[torch.Generator | None] | None = None, ) -> torch.Tensor: """Compute channel logits, apply repetition penalty, sample, store.""" channel_logits = audio_lm_heads[channel_index](local_hidden).float() @@ -254,6 +257,7 @@ class MossTTSLocalDepthTransformer(nn.Module): top_p, do_sample, generator=generator, + generators=generators, ) codes[:, channel_index] = channel_token return channel_token @@ -277,6 +281,7 @@ class MossTTSLocalDepthTransformer(nn.Module): repetition_penalty: float = 1.0, history_per_codebook: list[list[int]] | None = None, generator: torch.Generator | None = None, + generators: Sequence[torch.Generator | None] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Generate one audio frame for batch B. @@ -319,6 +324,7 @@ class MossTTSLocalDepthTransformer(nn.Module): text_top_p, do_sample, generator=generator, + generators=generators, ) should_continue = binary_choice.eq(0) import os as _os @@ -344,6 +350,7 @@ class MossTTSLocalDepthTransformer(nn.Module): top_p=top_p, do_sample=do_sample, generator=generator, + generators=generators, ) if channel_index + 1 < n_vq: diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index 249289dc3..bc4a5b716 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -30,6 +30,7 @@ from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import ( ) from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import ( MossTTSRealtimeLocalTransformer, + _normalize_generators, ) from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import ( MossTTSLocalDepthTransformer, @@ -1355,6 +1356,7 @@ class MossTTSLocalTalkerForGeneration(nn.Module): self._stacked_audio_emb_w: torch.Tensor | None = None self.mtp_hidden_size = hidden_size self.talker_mtp_graph_safe = not current_omni_platform.is_npu() + self.talker_mtp_accepts_per_row_generators = True self.talker_mtp_output_key = ("audio_codes", "current") # ``make_omni_output`` keeps code rows fixed-shape and performs all # state updates eagerly, so the runner can safely pack/snapshot them @@ -1621,9 +1623,14 @@ class MossTTSLocalTalkerForGeneration(nn.Module): top_k: int | None = None, top_p: float | None = None, generator: torch.Generator | None = None, + generators: list[torch.Generator | None] | None = None, **_: Any, ) -> tuple[torch.Tensor, torch.Tensor]: bsz = int(input_embeds.shape[0]) + # One generator per row, or the extra rows would silently sample from + # the global RNG. Validate before the n_vq-step depth loop so a wrong + # generator count fails on entry. + generators = _normalize_generators(generators, bsz) input_embeds_out = input_embeds.reshape(bsz, -1) last_talker_hidden = last_talker_hidden.reshape(bsz, -1).to( device=input_embeds.device, @@ -1652,6 +1659,7 @@ class MossTTSLocalTalkerForGeneration(nn.Module): repetition_penalty=1.0, history_per_codebook=None, generator=generator, + generators=generators, ) new_codes = new_codes.to(device=input_embeds.device, dtype=torch.long) emit_mask = active_mask & should_continue_t.reshape(bsz, 1)