refactor: clarify seeded Qwen3-TTS MTP generator dispatch

Signed-off-by: Sy03 <1370724210@qq.com>

tests/worker/test_omni_gpu_model_runner.py

@@ -171,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)
@@ -255,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):
@@ -269,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
 
@@ -280,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={
@@ -314,7 +350,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
     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)
@@ -329,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):
@@ -348,13 +385,13 @@ 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_batches_seeded_requests_for_opted_in_models(monkeypatch):
-    """Models with talker_mtp_accepts_per_row_generators get one batched call (#4883)."""
+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)
@@ -369,10 +406,7 @@ def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypa
         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,
-    )
+    _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):
@@ -386,8 +420,9 @@ def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypa
 
     # 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
+    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]
@@ -404,6 +439,62 @@ def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypa
     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/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -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,13 +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")
-        # 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"))
 
         if get_pp_group().is_last_rank:

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,62 +1897,22 @@ 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
-
-        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.
-            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
+        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"),
@@ -1881,11 +1920,10 @@ class OmniGPUModelRunner(GPUModelRunner):
             "top_k": subtalker_params.get("top_k"),
             "top_p": subtalker_params.get("top_p"),
         }
-        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 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"] = [