[Bugfix][Qwen3-TTS] Keep seeded residual MTP sampling batched (#4889)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -339,6 +339,91 @@ class TestCodePredictorDtypeAlignment:
         assert not torch.equal(first[:, 1:], different[:, 1:])
 
 
+class TestCodePredictorPerRowGenerators:
+    """Seeded requests must stay deterministic inside a multi-row batch (#4883)."""
+
+    def _make_predictor(self, mocker: MockerFixture, loaded_target_classes):
+        _, _, code_predictor_wrapper, _, _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+        cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+        vllm_config = _make_vllm_config(mocker, max_num_seqs=4)
+        predictor = code_predictor_wrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+        predictor._wrapper_config.use_cuda_graphs = False
+        return predictor, talker_config
+
+    def test_multinomial_per_row_matches_single_row_draws(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        multinomial = common_mod.CodePredictorWrapper._multinomial
+        torch.manual_seed(7)
+        probs = torch.softmax(torch.randn(3, 64), dim=-1)
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=probs.device)
+            generator.manual_seed(seed)
+            return generator
+
+        batched = multinomial(probs, None, [seeded(11), None, seeded(22)])
+        assert batched.shape == (3, 1)
+        # Each seeded row must reproduce a standalone draw from the same seed,
+        # independent of what the other rows in the batch consume.
+        assert torch.equal(batched[0:1], torch.multinomial(probs[0:1], num_samples=1, generator=seeded(11)))
+        assert torch.equal(batched[2:3], torch.multinomial(probs[2:3], num_samples=1, generator=seeded(22)))
+
+    def test_forward_per_row_generators_are_row_independent(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        torch.manual_seed(123)
+        row_embed = torch.randn(1, hidden)
+        row_hidden = torch.randn(1, hidden)
+        # Three identical rows: two share a seed, one differs.
+        layer0_code = torch.zeros(3, dtype=torch.long)
+        layer0_embed = row_embed.expand(3, hidden).contiguous()
+        last_talker_hidden = row_hidden.expand(3, hidden).contiguous()
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=layer0_code.device)
+            generator.manual_seed(seed)
+            return generator
+
+        def run(generators):
+            return predictor(
+                layer0_code=layer0_code,
+                layer0_embed=layer0_embed,
+                last_talker_hidden=last_talker_hidden,
+                do_sample=True,
+                temperature=0.9,
+                top_k=50,
+                top_p=1.0,
+                generators=generators,
+            )
+
+        first = run([seeded(1234), seeded(1234), seeded(4321)])
+        second = run([seeded(1234), seeded(1234), seeded(4321)])
+
+        # Same per-row seeds -> the whole batch reproduces across calls.
+        assert torch.equal(first, second)
+        # Identical inputs + identical seeds -> identical rows within a call.
+        assert torch.equal(first[0], first[1])
+        # A different seed on the same inputs must diverge in the residual layers.
+        assert not torch.equal(first[0, 1:], first[2, 1:])
+
+    def test_forward_rejects_mismatched_generators_length(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        with pytest.raises(ValueError, match="one entry per row"):
+            predictor(
+                layer0_code=torch.zeros(2, dtype=torch.long),
+                layer0_embed=torch.randn(2, hidden),
+                last_talker_hidden=torch.randn(2, hidden),
+                generators=[torch.Generator()],
+            )
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 

tests/worker/test_omni_gpu_model_runner.py

@@ -81,6 +81,7 @@ class CaptureTalkerMTP(torch.nn.Module):
         top_k=None,
         top_p=None,
         generator=None,
+        generators=None,
     ):
         self.calls.append(
             {
@@ -90,6 +91,7 @@ class CaptureTalkerMTP(torch.nn.Module):
                 "top_k": top_k,
                 "top_p": top_p,
                 "generator": generator,
+                "generators": generators,
             }
         )
         codes = torch.zeros((req_embeds.shape[0], 1), dtype=torch.int64)
@@ -306,6 +308,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
             "top_k": 9,
             "top_p": 0.55,
             "generator": runner.talker_mtp.calls[0]["generator"],
+            "generators": None,
         }
     ]
     assert runner.talker_mtp.calls[0]["generator"] is not None
@@ -350,6 +353,57 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
     assert torch.equal(runner.talker_mtp_inputs_embeds.gpu, saved_embeds)
 
 
+def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypatch):
+    """Models with talker_mtp_accepts_per_row_generators get one batched call (#4883)."""
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+    runner.requests["r1"].sampling_params = SimpleNamespace(
+        seed=11,
+        extra_args={"tts_local_seed": 11},
+    )
+    runner.requests["r2"].sampling_params = SimpleNamespace(
+        seed=22,
+        extra_args={"tts_local_seed": 22},
+    )
+    runner.talker_mtp = CaptureTalkerMTP()
+    runner.model = SimpleNamespace(
+        talker_mtp_output_key=("codes", "audio"),
+        talker_mtp_accepts_per_row_generators=True,
+    )
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
+
+    def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
+        batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
+        return (False, batch_desc, None, None, None)
+
+    monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
+
+    inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+
+    # One batched call with distinct per-row generators, not two scalar calls.
+    assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
+    row_generators = runner.talker_mtp.calls[0]["generators"]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    assert len(row_generators) == 2
+    assert all(generator is not None for generator in row_generators)
+    assert row_generators[0] is not row_generators[1]
+
+    # The per-request generator stream persists across steps...
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+    assert runner.talker_mtp.calls[1]["generators"][0] is row_generators[0]
+    assert runner.talker_mtp.calls[1]["generators"][1] is row_generators[1]
+
+    # ...and is evicted once its request finishes.
+    del runner.requests["r2"]
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
+    assert set(runner._talker_mtp_generators) == {"r1"}
+    assert runner.talker_mtp.calls[2]["generator"] is row_generators[0]
+
+
 def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):
     """Validate that _update_intermediate_buffer writes to model_intermediate_buffer
     (forward path) and mirrors to additional_information_cpu setattr (backward compat)."""

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -14,7 +14,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models.
 from __future__ import annotations
 
 import dataclasses
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 
 import torch
 import torch.nn as nn
@@ -742,6 +742,27 @@ class CodePredictorWrapper(nn.Module):
     #  Forward -- re-prefill + inline sampling
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _multinomial(
+        probs: torch.Tensor,
+        generator: torch.Generator | None,
+        generators: Sequence[torch.Generator | None] | None,
+    ) -> torch.Tensor:
+        """Sample one code per row, optionally with per-row generators.
+
+        Per-row generators keep explicitly-seeded requests deterministic in a
+        multi-row batch: each row consumes draws only from its own generator,
+        so the transformer forward can stay batched (#4883).
+        """
+        if generators is None:
+            return torch.multinomial(probs, num_samples=1, generator=generator)
+        return torch.cat(
+            [
+                torch.multinomial(probs[row : row + 1], num_samples=1, generator=row_generator)
+                for row, row_generator in enumerate(generators)
+            ]
+        )
+
     @torch.inference_mode()
     def forward(
         self,
@@ -753,9 +774,12 @@ class CodePredictorWrapper(nn.Module):
         top_k: int = 50,
         top_p: float = 1.0,
         generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """Predict residual codebooks 1..G-1 autoregressively via re-prefill."""
         bsz = int(layer0_code.shape[0])
+        if generators is not None and len(generators) != bsz:
+            raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
         num_groups = self._num_groups
         device = layer0_code.device
 
@@ -846,7 +870,7 @@ class CodePredictorWrapper(nn.Module):
                     sorted_logits[remove_mask] = float("-inf")
                     logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
                 probs = F.softmax(logits, dim=-1, dtype=torch.float32)
-                code = torch.multinomial(probs, num_samples=1, generator=generator)
+                code = self._multinomial(probs, generator, generators)
             else:
                 # "per_call" mode: temperature-scaled + top-k
                 if use_sampling:
@@ -855,7 +879,7 @@ class CodePredictorWrapper(nn.Module):
                         topk_vals, _ = scaled.topk(top_k, dim=-1)
                         scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
                     probs = F.softmax(scaled, dim=-1, dtype=torch.float32)
-                    code = torch.multinomial(probs, num_samples=1, generator=generator)
+                    code = self._multinomial(probs, generator, generators)
                 else:
                     code = logits.argmax(dim=-1, keepdim=True)
 

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any
 
 import numpy as np
@@ -318,6 +318,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # OmniGPUModelRunner will store talker_mtp output under this key in
         # per-request additional_information.
         self.talker_mtp_output_key = ("codes", "audio")
+        # talker_mtp samples with per-row generators, so explicitly-seeded
+        # requests stay batched instead of one scalar forward per row (#4883).
+        # Only valid while talker_mtp receives the unpadded active batch (this
+        # talker is not graph-wrapped); a padded batch would need the runner to
+        # pad the generators list as well.
+        self.talker_mtp_accepts_per_row_generators = True
 
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
@@ -1056,6 +1062,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         top_k: int | None = None,
         top_p: float | None = None,
         generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
         **kwargs: Any,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """GPU fast-path used by OmniGPUModelRunner to predict residual codebooks (1..Q-1).
@@ -1095,6 +1102,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             top_k=top_k,
             top_p=top_p,
             generator=generator,
+            generators=generators,
         )  # [B, Q]
 
         # Map invalid layer-0 ids (e.g. EOS) to PAD=0 so SpeechTokenizer sees only real codes.

vllm_omni/worker/gpu_model_runner.py

@@ -1826,7 +1826,33 @@ class OmniGPUModelRunner(GPUModelRunner):
                 seed = extra_args.get("tts_local_seed")
             return int(seed) if seed is not None else None
 
-        if decode_batch_size > 1 and any(_explicit_talker_seed(req_id) is not None for req_id in decode_req_ids):
+        def _row_generator(req_id: str) -> torch.Generator | None:
+            seed = _explicit_talker_seed(req_id)
+            if seed is None:
+                return None
+            cache = getattr(self, "_talker_mtp_generators", None)
+            if cache is None:
+                cache = {}
+                self._talker_mtp_generators = cache
+            generator = cache.get(req_id)
+            if generator is None or generator.device != req_input_ids.device:
+                generator = torch.Generator(device=req_input_ids.device)
+                generator.manual_seed(seed)
+                cache[req_id] = generator
+            return generator
+
+        row_generators = [_row_generator(req_id) for req_id in decode_req_ids]
+        cache = getattr(self, "_talker_mtp_generators", None)
+        if cache:
+            # Generators live as long as their request; drop finished ones.
+            for stale_id in [rid for rid in cache if rid not in self.requests]:
+                del cache[stale_id]
+
+        if (
+            decode_batch_size > 1
+            and any(generator is not None for generator in row_generators)
+            and not getattr(self.model, "talker_mtp_accepts_per_row_generators", False)
+        ):
             # A torch.Generator is a single stream. Using one generator for a
             # multi-row batch would make explicitly-seeded requests depend on
             # other rows in the same scheduler step, so keep that path scalar.
@@ -1849,28 +1875,17 @@ class OmniGPUModelRunner(GPUModelRunner):
                 self.text_step.gpu[:decode_batch_size].copy_(saved_text)
             return
 
-        generator = None
-        if decode_req_ids:
-            first_req_id = decode_req_ids[0]
-            seed = _explicit_talker_seed(first_req_id)
-            if seed is not None:
-                generators = getattr(self, "_talker_mtp_generators", None)
-                if generators is None:
-                    generators = {}
-                    self._talker_mtp_generators = generators
-                generator = generators.get(first_req_id)
-                if generator is None or generator.device != req_input_ids.device:
-                    generator = torch.Generator(device=req_input_ids.device)
-                    generator.manual_seed(int(seed))
-                    generators[first_req_id] = generator
         talker_kwargs = {
             "do_sample": subtalker_params.get("do_sample"),
             "temperature": subtalker_params.get("temperature"),
             "top_k": subtalker_params.get("top_k"),
             "top_p": subtalker_params.get("top_p"),
         }
-        if generator is not None:
-            talker_kwargs["generator"] = generator
+        if decode_batch_size == 1:
+            if row_generators[0] is not None:
+                talker_kwargs["generator"] = row_generators[0]
+        elif any(generator is not None for generator in row_generators):
+            talker_kwargs["generators"] = row_generators
         if getattr(self.model, "talker_mtp_accepts_req_infos", False):
             talker_kwargs["req_ids"] = decode_req_ids
             talker_kwargs["req_infos"] = [