refactor: clarify seeded Qwen3-TTS MTP generator dispatch

Signed-off-by: Sy03 <1370724210@qq.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)
@@ -169,6 +171,12 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
     return runner
 
 
+def _enable_talker_seed(runner, *, accepts_generators=False):
+    runner.model.talker_mtp_seed_extra_arg = "tts_local_seed"
+    if accepts_generators:
+        runner.model.talker_mtp_accepts_generators = True
+
+
 def _make_runner_for_mimo(req_id="r_mimo"):
     """Create a minimal runner with MiMoAudio-like model and request state."""
     runner = object.__new__(OmniGPUModelRunner)
@@ -253,6 +261,7 @@ def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(seed=42)
     runner.talker_mtp = CaptureTalkerMTP()
+    _enable_talker_seed(runner)
     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):
@@ -267,6 +276,34 @@ def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker
     assert runner.talker_mtp.calls[0]["generator"] is None
 
 
+def test_talker_mtp_forward_ignores_tts_local_seed_without_model_seed_key(monkeypatch):
+    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=42,
+        extra_args={"tts_local_seed": 42},
+    )
+    runner.talker_mtp = CaptureTalkerMTP()
+    runner.model.talker_mtp_accepts_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)
+
+    assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    assert runner.talker_mtp.calls[0]["generators"] is None
+
+
 def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
@@ -278,6 +315,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
         extra_args={"tts_local_seed": 42},
     )
     runner.talker_mtp = CaptureTalkerMTP()
+    _enable_talker_seed(runner)
     runner.vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(
             subtalker_sampling_params={
@@ -306,12 +344,13 @@ 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
 
 
-def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
+def test_talker_mtp_forward_keeps_unmarked_models_scalar_for_explicit_seed(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
     monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
@@ -326,6 +365,7 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
         extra_args={"tts_local_seed": 22},
     )
     runner.talker_mtp = CaptureTalkerMTP()
+    _enable_talker_seed(runner)
     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):
@@ -345,11 +385,116 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
 
     assert [call["batch_size"] for call in runner.talker_mtp.calls] == [1, 1]
     assert all(call["generator"] is not None for call in runner.talker_mtp.calls)
+    assert all(call["generators"] is None for call in runner.talker_mtp.calls)
     assert runner.talker_mtp.calls[0]["generator"] is not runner.talker_mtp.calls[1]["generator"]
     assert torch.equal(runner.talker_mtp_input_ids.gpu, saved_input_ids)
     assert torch.equal(runner.talker_mtp_inputs_embeds.gpu, saved_embeds)
 
 
+def test_talker_mtp_forward_keeps_explicit_seeded_requests_batched(monkeypatch):
+    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()
+    _enable_talker_seed(runner, accepts_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]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    row_generators = runner.talker_mtp.calls[0]["generators"]
+    assert row_generators is not 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_talker_mtp_forward_batches_mixed_explicit_seed_requests(monkeypatch):
+    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=("seeded", "unseeded"), hidden_size=4)
+    runner.requests["seeded"].sampling_params = SimpleNamespace(
+        seed=11,
+        extra_args={"tts_local_seed": 11},
+    )
+    runner.talker_mtp = CaptureTalkerMTP()
+    _enable_talker_seed(runner, accepts_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, ["seeded", "unseeded"], inputs_embeds)
+
+    assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    generators = runner.talker_mtp.calls[0]["generators"]
+    assert generators is not None
+    assert len(generators) == 2
+    assert generators[0] is not None
+    assert generators[1] is None
+
+
+def test_talker_mtp_forward_does_not_pass_generators_without_explicit_seed(monkeypatch):
+    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.talker_mtp = CaptureTalkerMTP()
+    _enable_talker_seed(runner, accepts_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)
+
+    assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    assert runner.talker_mtp.calls[0]["generators"] is None
+
+
 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/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -2235,13 +2235,8 @@ class HunyuanImage3Model(nn.Module):
             # processed with quantization, LoRA, fine-tuning, etc.
             if self.config.tie_word_embeddings and "lm_head.weight" in name:
                 continue
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                # Loading kv cache scales for compressed-tensors quantization
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                continue
+            # KV-cache scales are renamed via maybe_remap_kv_scale_name below;
+            # quant_config.get_cache_scale was removed in vLLM v0.23.0 (see #4810).
 
             is_found = False
             for param_name, weight_name, shard_id in stacked_params_mapping:

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
@@ -266,6 +266,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
     """vLLM-AR talker: step-wise layer-0 codec decoding.
     Predicts residual codebooks (1..Q-1) into `audio_codes` and streams text via `tailing_text_hidden`."""
 
+    talker_mtp_accepts_generators = True
+    talker_mtp_seed_extra_arg = "tts_local_seed"
+
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             # Talker backbone (Qwen3 decoder-only).
@@ -318,7 +321,6 @@ 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")
-
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
         if get_pp_group().is_last_rank:
@@ -1056,6 +1058,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 +1098,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

@@ -84,6 +84,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         self.omni_prefix_cache = None
         self._sampled_token_ids_cpu_override = None
         self._omni_query_start_loc_model_kwarg = False
+        self._talker_mtp_generators: dict[str, torch.Generator] = {}
 
     def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
         override_fn = self._sampled_token_ids_cpu_override
@@ -1785,6 +1786,84 @@ class OmniGPUModelRunner(GPUModelRunner):
             ec_connector_output,
         )
 
+    def _explicit_talker_seed(self, req_id: str) -> int | None:
+        seed_extra_arg = getattr(self.model, "talker_mtp_seed_extra_arg", None)
+        if not seed_extra_arg:
+            return None
+
+        sampling_params = getattr(self.requests[req_id], "sampling_params", None)
+        extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
+        if not isinstance(extra_args, dict):
+            return None
+
+        seed = extra_args.get(seed_extra_arg)
+        return int(seed) if seed is not None else None
+
+    def _resolve_talker_mtp_generators(
+        self,
+        decode_req_ids: list[str],
+        device: torch.device,
+    ) -> list[torch.Generator | None] | None:
+        """Resolve per-row generators aligned to ``decode_req_ids``.
+
+        ``row_generators[i]`` belongs to the batched input row for
+        ``decode_req_ids[i]``. Explicitly seeded rows use cached per-request
+        generators, while unseeded rows stay as ``None`` to keep global RNG
+        semantics. Return ``None`` when the whole batch is unseeded so callers
+        can use the normal batched path without sampling kwargs.
+        """
+        generator_cache = getattr(self, "_talker_mtp_generators", None)
+        if generator_cache is None:
+            generator_cache = {}
+            self._talker_mtp_generators = generator_cache
+
+        row_generators: list[torch.Generator | None] = []
+        any_seeded = False
+        for req_id in decode_req_ids:
+            seed = self._explicit_talker_seed(req_id)
+            if seed is None:
+                row_generators.append(None)
+                continue
+
+            generator = generator_cache.get(req_id)
+            if generator is None or generator.device != device:
+                generator = torch.Generator(device=device)
+                generator.manual_seed(seed)
+                generator_cache[req_id] = generator
+            row_generators.append(generator)
+            any_seeded = True
+
+        for stale_id in [rid for rid in generator_cache if rid not in self.requests]:
+            del generator_cache[stale_id]
+
+        return row_generators if any_seeded else None
+
+    def _talker_mtp_forward_row_by_row(
+        self,
+        decode_req_ids: list[str],
+        inputs_embeds: torch.Tensor,
+        start_offsets: list[int] | None = None,
+    ) -> None:
+        """Seeded fallback for models without batched per-row generators."""
+        decode_batch_size = len(decode_req_ids)
+        saved_input_ids = self.talker_mtp_input_ids.gpu[:decode_batch_size].clone()
+        saved_embeds = self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].clone()
+        saved_hidden = self.last_talker_hidden.gpu[:decode_batch_size].clone()
+        saved_text = self.text_step.gpu[:decode_batch_size].clone()
+        try:
+            for row, req_id in enumerate(decode_req_ids):
+                self.talker_mtp_input_ids.gpu[:1].copy_(saved_input_ids[row : row + 1])
+                self.talker_mtp_inputs_embeds.gpu[:1].copy_(saved_embeds[row : row + 1])
+                self.last_talker_hidden.gpu[:1].copy_(saved_hidden[row : row + 1])
+                self.text_step.gpu[:1].copy_(saved_text[row : row + 1])
+                row_offsets = None if start_offsets is None else [start_offsets[row]]
+                self._talker_mtp_forward([req_id], inputs_embeds, row_offsets)
+        finally:
+            self.talker_mtp_input_ids.gpu[:decode_batch_size].copy_(saved_input_ids)
+            self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].copy_(saved_embeds)
+            self.last_talker_hidden.gpu[:decode_batch_size].copy_(saved_hidden)
+            self.text_step.gpu[:decode_batch_size].copy_(saved_text)
+
     def _talker_mtp_forward(
         self,
         decode_req_ids: list[str],
@@ -1818,51 +1897,23 @@ class OmniGPUModelRunner(GPUModelRunner):
         if not isinstance(subtalker_params, dict):
             subtalker_params = {}
 
-        def _explicit_talker_seed(req_id: str) -> int | None:
-            sampling_params = getattr(self.requests[req_id], "sampling_params", None)
-            extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
-            seed = None
-            if isinstance(extra_args, dict):
-                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):
-            # 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.
-            saved_input_ids = self.talker_mtp_input_ids.gpu[:decode_batch_size].clone()
-            saved_embeds = self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].clone()
-            saved_hidden = self.last_talker_hidden.gpu[:decode_batch_size].clone()
-            saved_text = self.text_step.gpu[:decode_batch_size].clone()
-            try:
-                for row, req_id in enumerate(decode_req_ids):
-                    self.talker_mtp_input_ids.gpu[:1].copy_(saved_input_ids[row : row + 1])
-                    self.talker_mtp_inputs_embeds.gpu[:1].copy_(saved_embeds[row : row + 1])
-                    self.last_talker_hidden.gpu[:1].copy_(saved_hidden[row : row + 1])
-                    self.text_step.gpu[:1].copy_(saved_text[row : row + 1])
-                    row_offsets = None if start_offsets is None else [start_offsets[row]]
-                    self._talker_mtp_forward([req_id], inputs_embeds, row_offsets)
-            finally:
-                self.talker_mtp_input_ids.gpu[:decode_batch_size].copy_(saved_input_ids)
-                self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].copy_(saved_embeds)
-                self.last_talker_hidden.gpu[:decode_batch_size].copy_(saved_hidden)
-                self.text_step.gpu[:decode_batch_size].copy_(saved_text)
-            return
-
+        row_generators = self._resolve_talker_mtp_generators(
+            decode_req_ids,
+            req_input_ids.device,
+        )
         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
+        generators = None
+        if row_generators is not None:
+            if decode_batch_size == 1:
+                generator = row_generators[0]
+            elif getattr(self.model, "talker_mtp_accepts_generators", False):
+                # Unseeded rows intentionally keep the default global RNG semantics,
+                # represented by None, even when batched with explicitly seeded rows.
+                generators = row_generators
+            else:
+                self._talker_mtp_forward_row_by_row(decode_req_ids, inputs_embeds, start_offsets)
+                return
+
         talker_kwargs = {
             "do_sample": subtalker_params.get("do_sample"),
             "temperature": subtalker_params.get("temperature"),
@@ -1871,6 +1922,8 @@ class OmniGPUModelRunner(GPUModelRunner):
         }
         if generator is not None:
             talker_kwargs["generator"] = generator
+        if generators is not None:
+            talker_kwargs["generators"] = generators
         if getattr(self.model, "talker_mtp_accepts_req_infos", False):
             talker_kwargs["req_ids"] = decode_req_ids
             talker_kwargs["req_infos"] = [