[Perf][TTS] Add per-row generator support to MOSS-TTS talker to eliminate serial fallback for seeded batched requests (#7922)

Signed-off-by: jingchengtian <jingchengtian@users.noreply.github.com>
Signed-off-by: jingchengtian <tjc1995@126.com>
Co-authored-by: jingchengtian <jingchengtian@users.noreply.github.com>

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])

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

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:

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)