From 95934535f31926ee41f455e2117c0ee2869684f2 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Sun, 5 Jul 2026 22:49:43 +0800 Subject: [PATCH] refactor: clarify seeded Qwen3-TTS MTP generator dispatch Signed-off-by: Sy03 <1370724210@qq.com> --- tests/worker/test_omni_gpu_model_runner.py | 107 +++++++++++- .../models/qwen3_tts/qwen3_tts_talker.py | 10 +- vllm_omni/worker/gpu_model_runner.py | 160 +++++++++++------- 3 files changed, 201 insertions(+), 76 deletions(-) diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index 48c0bc6de..3d0b5d349 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/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).""" 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 347c76dfa..c40258995 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 @@ -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: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index e922146db..b4ff3ecd0 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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"] = [