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 87281af82..b73c15fc7 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 @@ -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.""" diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index e68dd6c43..48c0bc6de 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/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) @@ -306,6 +308,7 @@ 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 @@ -350,6 +353,57 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch): 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).""" + 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() + runner.model = SimpleNamespace( + talker_mtp_output_key=("codes", "audio"), + talker_mtp_accepts_per_row_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] + row_generators = runner.talker_mtp.calls[0]["generators"] + assert runner.talker_mtp.calls[0]["generator"] is 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_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/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index eca5addba..c8591c32e 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/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) 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 c1c7f4264..347c76dfa 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 @@ -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 @@ -318,6 +318,12 @@ 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")) @@ -1056,6 +1062,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 +1102,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. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index e40c220b8..e922146db 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1826,7 +1826,33 @@ class OmniGPUModelRunner(GPUModelRunner): 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): + 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. @@ -1849,28 +1875,17 @@ class OmniGPUModelRunner(GPUModelRunner): self.text_step.gpu[:decode_batch_size].copy_(saved_text) return - 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 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 generator is not None: - talker_kwargs["generator"] = generator + 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 getattr(self.model, "talker_mtp_accepts_req_infos", False): talker_kwargs["req_ids"] = decode_req_ids talker_kwargs["req_infos"] = [