From 207da92c070f6ad500a1d1fc9728cbb5a196d222 Mon Sep 17 00:00:00 2001 From: jingchengtian Date: Mon, 21 Sep 2026 06:38:29 +0000 Subject: [PATCH 1/3] [Perf][TTS] Add per-row generator support to MOSS-TTS talker to eliminate serial fallback for seeded batched requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When B>1 seeded requests are batched, the runner serially processes each row (B×12 serial depth-transformer forward passes) unless the model sets talker_mtp_accepts_per_row_generators. This PR sets that flag and threads generators through _sample_token → generate_frame → talker_mtp, matching the Qwen3-TTS pattern (qwen3_tts_talker.py:351). Each row uses its own torch.Generator for multinomial, preserving per-row reproducibility. The expensive topk/softmax/nucleus filtering stays batched. AI assistance: Used opencode to draft the code, benchmark scripts, and PR description. I reviewed all changes and verified benchmark results from actual NPU runs. Validation: single-module + E2E benchmarks at B/C=1,8,16,32,64. Signed-off-by: jingchengtian Signed-off-by: jingchengtian --- .../moss_tts/modeling_moss_tts_local.py | 24 ++++++++++++++++--- .../moss_tts/modeling_moss_tts_local_depth.py | 7 ++++++ .../moss_tts/modeling_moss_tts_talker.py | 3 +++ 3 files changed, 31 insertions(+), 3 deletions(-) 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 c9ce11a935a..559f3d05c54 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 @@ -23,7 +23,7 @@ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence import torch import torch.nn as nn @@ -145,6 +145,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 +154,10 @@ 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. """ if not do_sample or temperature <= 0: return logits.argmax(dim=-1) @@ -181,8 +186,21 @@ 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]) + + if generators is not None and any(g is not None for g in generators): + B = probs.shape[0] + rows = [] + for row in range(B): + gen = generators[row] if row < len(generators) else None + row_probs = probs[row : row + 1] + flat = row_probs.reshape(-1, row_probs.shape[-1]) + sampled = torch.multinomial(flat, num_samples=1, generator=gen) + 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 a7e7606a32a..c3f226fd189 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 @@ 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 @@ def _sample_channel( 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 @@ def _sample_channel( top_p, do_sample, generator=generator, + generators=generators, ) codes[:, channel_index] = channel_token return channel_token @@ -277,6 +281,7 @@ def generate_frame( 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 @@ def generate_frame( text_top_p, do_sample, generator=generator, + generators=generators, ) should_continue = binary_choice.eq(0) import os as _os @@ -344,6 +350,7 @@ def generate_frame( 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 249289dc36d..0c8afe1c2b9 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 @@ -1355,6 +1355,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: 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,6 +1622,7 @@ def talker_mtp( 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]) @@ -1652,6 +1654,7 @@ def talker_mtp( 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) From 1327dffb6ae722152692595230475bdbb75b0f69 Mon Sep 17 00:00:00 2001 From: jingchengtian Date: Tue, 29 Sep 2026 03:29:02 +0000 Subject: [PATCH 2/3] [Perf][TTS] Validate per-row generator count in MOSS-TTS sampler A ``generators`` list shorter than the batch silently sampled the extra rows from the global RNG, so a mis-sized batch would quietly lose per-row reproducibility. Validate the length up front instead, mirroring ``Qwen3CodePredictor._normalize_generators``: ``_sample_token`` raises ``ValueError`` on a mismatch and ``talker_mtp`` fails on entry (before the n_vq-step depth loop) when the batch does not match. The all-``None`` case still falls through to the single batched ``multinomial``, so unseeded and B=1 requests keep the original path. Tests: new tests/model_executor/models/moss_tts/test_per_row_generators.py (CPU, 6 cases) covers the length check plus per-row reproducibility, batch composition independence and the unseeded fall-through. Validated on Ascend 910B2C with MOSS-TTS-Local-Transformer-v1.5: seeded batches of 4/8/16 return 200 with valid audio, C=4 1.61s vs 2.40s and C=16 2.81s vs 7.53s against the serial fallback, and a batched generate_frame reproduces per-row codes bit-identically to running each row alone. Signed-off-by: jingchengtian Signed-off-by: jingchengtian --- .../moss_tts/test_per_row_generators.py | 120 ++++++++++++++++++ .../moss_tts/modeling_moss_tts_local.py | 29 ++++- .../moss_tts/modeling_moss_tts_talker.py | 5 + 3 files changed, 149 insertions(+), 5 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 00000000000..ec1d8f047c6 --- /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 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 559f3d05c54..a617fbe7c61 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 @@ -138,6 +138,24 @@ def generate_frame( 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, @@ -158,6 +176,8 @@ def _sample_token( 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) @@ -186,15 +206,14 @@ def _sample_token( logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_indices, logits) probs = F.softmax(logits, dim=-1) - - if generators is not None and any(g is not None for g in generators): - B = probs.shape[0] + 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): - gen = generators[row] if row < len(generators) else None row_probs = probs[row : row + 1] flat = row_probs.reshape(-1, row_probs.shape[-1]) - sampled = torch.multinomial(flat, num_samples=1, generator=gen) + 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: 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 0c8afe1c2b9..c08c8e363d9 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.modeling_moss_tts_local import ( MossTTSRealtimeLocalTransformer, + _normalize_generators, ) from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import ( MossTTSLocalDepthTransformer, @@ -1626,6 +1627,10 @@ def talker_mtp( **_: 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 + # mis-sized batch 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, From f34ae58565d2b65a9ca44692a5d077a9e3f7061d Mon Sep 17 00:00:00 2001 From: jingchengtian Date: Tue, 29 Sep 2026 04:12:58 +0000 Subject: [PATCH 3/3] [Bugfix][TTS] Fix pre-commit failures in the per-row generator validation Two hooks rejected the previous commit: - check-spdx-header: modeling_moss_tts_local.py carried no SPDX header and the new test file copied the stale "vLLM project" copyright line. The hook rewrites both to the vLLM-Omni lines. - typos: "mis-sized" in the talker_mtp comment is not a word it knows. Comment and header only, no functional change. Signed-off-by: jingchengtian --- .../model_executor/models/moss_tts/test_per_row_generators.py | 2 +- .../model_executor/models/moss_tts/modeling_moss_tts_local.py | 3 +++ .../models/moss_tts/modeling_moss_tts_talker.py | 4 ++-- 3 files changed, 6 insertions(+), 3 deletions(-) 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 index ec1d8f047c6..4719ec8cb8b 100644 --- a/tests/model_executor/models/moss_tts/test_per_row_generators.py +++ b/tests/model_executor/models/moss_tts/test_per_row_generators.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 """CPU regression tests for MOSS-TTS per-row talker generators.""" import pytest 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 a617fbe7c61..f1d23a7deb8 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"). 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 c08c8e363d9..bc4a5b7163e 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 @@ -1628,8 +1628,8 @@ def talker_mtp( ) -> 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 - # mis-sized batch fails on entry. + # 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(