[Bugfix] Propagate seed to Qwen3-TTS Fast AR sampler (#3350)

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

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -268,6 +268,68 @@ class TestCodePredictorDtypeAlignment:
         assert result.shape == (bsz, num_groups)
         assert result.dtype == torch.long
 
+    def test_forward_generator_controls_sampling(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        _, _, 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=2)
+
+        predictor = code_predictor_wrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+        predictor._wrapper_config.use_cuda_graphs = False
+
+        bsz = 1
+        hidden = talker_config.hidden_size
+        torch.manual_seed(123)
+        layer0_code = torch.zeros(bsz, dtype=torch.long)
+        layer0_embed = torch.randn(bsz, hidden)
+        last_talker_hidden = torch.randn(bsz, hidden)
+
+        first_generator = torch.Generator(device=layer0_code.device)
+        first_generator.manual_seed(1234)
+        second_generator = torch.Generator(device=layer0_code.device)
+        second_generator.manual_seed(1234)
+        different_generator = torch.Generator(device=layer0_code.device)
+        different_generator.manual_seed(4321)
+
+        first = 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,
+            generator=first_generator,
+        )
+        second = 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,
+            generator=second_generator,
+        )
+        different = 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,
+            generator=different_generator,
+        )
+
+        assert torch.equal(first, second)
+        assert not torch.equal(first[:, 1:], different[:, 1:])
+
 
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""

tests/worker/test_omni_gpu_model_runner.py

@@ -78,6 +78,7 @@ class CaptureTalkerMTP(torch.nn.Module):
         temperature=None,
         top_k=None,
         top_p=None,
+        generator=None,
     ):
         self.calls.append(
             {
@@ -85,6 +86,7 @@ class CaptureTalkerMTP(torch.nn.Module):
                 "temperature": temperature,
                 "top_k": top_k,
                 "top_p": top_p,
+                "generator": generator,
             }
         )
         codes = torch.zeros((req_embeds.shape[0], 1), dtype=torch.int64)
@@ -208,12 +210,38 @@ def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):
     assert torch.allclose(inputs_embeds, before)
 
 
+def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker(monkeypatch):
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("r1",), hidden_size=4)
+    runner.requests["r1"].sampling_params = SimpleNamespace(seed=42)
+    runner.talker_mtp = CaptureTalkerMTP()
+    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((2, 4), dtype=torch.float32)
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
+
+    assert runner.talker_mtp.calls[0]["generator"] 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
 
     monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
+    runner.requests["r1"].sampling_params = SimpleNamespace(
+        seed=42,
+        extra_args={"qwen3_tts_request_seed": 42},
+    )
     runner.talker_mtp = CaptureTalkerMTP()
     runner.vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(
@@ -241,8 +269,10 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
             "temperature": 0.2,
             "top_k": 9,
             "top_p": 0.55,
+            "generator": runner.talker_mtp.calls[0]["generator"],
         }
     ]
+    assert runner.talker_mtp.calls[0]["generator"] is not None
 
 
 def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):

vllm_omni/entrypoints/openai/serving_speech.py

@@ -2120,21 +2120,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     request.max_new_tokens,
                 )
 
-        # Propagate per-request seed to sampling params so both Slow AR
-        # and Fast AR produce deterministic output for the same seed.
         if request.seed is not None and sampling_params_list:
-            if not self._is_fish_speech:
-                logger.warning(
-                    "seed=%d requested but deterministic Fast AR seeding is "
-                    "only implemented for Fish Speech; other TTS models will "
-                    "use the seed for the main AR sampler only.",
-                    request.seed,
-                )
             if sampling_params_list is self.engine_client.default_sampling_params_list:
                 import copy
 
                 sampling_params_list = copy.deepcopy(sampling_params_list)
             sampling_params_list[0].seed = request.seed
+            if self._tts_model_type == "qwen3_tts":
+                if sampling_params_list[0].extra_args is None:
+                    sampling_params_list[0].extra_args = {}
+                sampling_params_list[0].extra_args["qwen3_tts_request_seed"] = request.seed
 
         generator = self.engine_client.generate(
             prompt=prompt,

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -619,6 +619,7 @@ class CodePredictorWrapper(nn.Module):
         temperature: float = 0.9,
         top_k: int = 50,
         top_p: float = 1.0,
+        generator: torch.Generator | 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])
@@ -703,7 +704,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)
+                code = torch.multinomial(probs, num_samples=1, generator=generator)
             else:
                 # "per_call" mode: temperature-scaled + top-k
                 if use_sampling:
@@ -712,7 +713,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)
+                    code = torch.multinomial(probs, num_samples=1, generator=generator)
                 else:
                     code = logits.argmax(dim=-1, keepdim=True)
 

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -1681,6 +1681,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         temperature: float | None = None,
         top_k: int | None = None,
         top_p: float | None = None,
+        generator: torch.Generator | None = None,
         **kwargs: Any,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """GPU fast-path used by OmniGPUModelRunner to predict residual codebooks (1..Q-1).
@@ -1718,6 +1719,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             temperature=temperature,
             top_k=top_k,
             top_p=top_p,
+            generator=generator,
         )  # [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

@@ -270,6 +270,8 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.num_prompt_logprobs.pop(req_id, None)
             if hasattr(self, "_downstream_payload_cache"):
                 self._downstream_payload_cache.pop(req_id, None)
+            if hasattr(self, "_talker_mtp_generators"):
+                self._talker_mtp_generators.pop(req_id, None)
 
         if hasattr(self, "late_interaction_runner"):
             self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids)
@@ -1384,35 +1386,41 @@ class OmniGPUModelRunner(GPUModelRunner):
         subtalker_params = getattr(self.vllm_config.model_config, "subtalker_sampling_params", None)
         if not isinstance(subtalker_params, dict):
             subtalker_params = {}
-        # Extract seed from the first request's sampling params for Fast AR
-        # determinism. NOTE: when batch_size > 1, all requests share the first
-        # request's seed. Per-request Fast AR seeding requires row-by-row
-        # torch.multinomial calls and is left as a follow-up optimisation.
-        _seed = None
+        generator = None
         if decode_req_ids:
-            _first_sp = getattr(self.requests[decode_req_ids[0]], "sampling_params", None)
-            if _first_sp is not None and getattr(_first_sp, "seed", None) is not None:
-                _seed = _first_sp.seed
-            # Warn when batched requests have different seeds.
-            if len(decode_req_ids) > 1 and _seed is not None:
-                _other_seeds = {
+            first_req_id = decode_req_ids[0]
+            first_sp = getattr(self.requests[first_req_id], "sampling_params", None)
+            extra_args = getattr(first_sp, "extra_args", None) if first_sp is not None else None
+            seed = extra_args.get("qwen3_tts_request_seed") if isinstance(extra_args, dict) else None
+            if len(decode_req_ids) > 1 and seed is not None:
+                other_seeds = {
                     getattr(getattr(self.requests[rid], "sampling_params", None), "seed", None)
                     for rid in decode_req_ids[1:]
                 }
-                if _other_seeds != {_seed}:
+                if other_seeds != {seed}:
                     logger.warning(
                         "Fast AR seed: batch has mixed seeds; using first request's seed=%d for all %d requests.",
-                        _seed,
+                        seed,
                         len(decode_req_ids),
                     )
+            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 _seed is not None:
-            talker_kwargs["seed"] = _seed
+        if generator is not None:
+            talker_kwargs["generator"] = generator
         with set_forward_context(
             None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
         ):