From 2b85474c2347d360df1e245ee6e6b628041536fb Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Wed, 6 May 2026 02:57:55 +0800 Subject: [PATCH] [Bugfix] Propagate seed to Qwen3-TTS Fast AR sampler (#3350) Signed-off-by: Sy03 <1370724210@qq.com> --- .../qwen3_tts/test_code_predictor_dtype.py | 62 +++++++++++++++++++ tests/worker/test_omni_gpu_model_runner.py | 30 +++++++++ .../entrypoints/openai/serving_speech.py | 13 ++-- .../models/common/qwen3_code_predictor.py | 5 +- .../models/qwen3_tts/qwen3_tts_talker.py | 2 + vllm_omni/worker/gpu_model_runner.py | 38 +++++++----- 6 files changed, 124 insertions(+), 26 deletions(-) diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 8798cb3ca..3c76f4210 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/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.""" diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index 9f908a7ad..c2f3f3622 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/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): diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index f6477bdca..9b9a7f205 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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, diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 70c96fb84..d36068d46 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/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) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index d9c454bd7..5703377da 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/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. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 0f7ed429e..f4eadedc0 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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 ):