Merge main into dev/vllm-align for vLLM 0.31 release

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/worker/test_gpu_ar_model_runner.py

@@ -15,6 +15,8 @@ import numpy as np
 import pytest
 import torch
 from vllm.sampling_params import SamplingParams
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.utils import CpuGpuBuffer
 from vllm.v1.worker import gpu_input_batch
 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 
@@ -103,6 +105,8 @@ def test_post_sample_talker_mtp_uses_current_sample_and_hidden() -> None:
     )
 
     assert received["req_ids"] == ["ready"]
+    assert isinstance(received["input_ids"], torch.Tensor)
+    assert isinstance(received["hidden_states"], torch.Tensor)
     assert received["input_ids"].tolist() == [101]
     assert received["hidden_states"].tolist() == [[1.0, 2.0]]
     assert received["req_infos"] == [{"duplex": {"data_plane": True}}]
@@ -139,6 +143,8 @@ def test_post_sample_talker_mtp_uses_gpu_token_with_async_scheduling() -> None:
     )
 
     assert received["req_ids"] == ["ready"]
+    assert isinstance(received["input_ids"], torch.Tensor)
+    assert isinstance(received["hidden_states"], torch.Tensor)
     assert received["input_ids"].tolist() == [101]
     assert received["hidden_states"].tolist() == [[1.0, 2.0]]
     assert received["req_infos"] == [{"duplex": {"data_plane": True}}]
@@ -275,7 +281,12 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
     )
     input_batch.sampling_metadata = input_batch._make_sampling_metadata()
 
-    received = []
+    received: list[SamplingMetadata] = []
+
+    def sample(logits: torch.Tensor, metadata: SamplingMetadata) -> str:
+        received.append(metadata)
+        return "model-sampler"
+
     runner = object.__new__(GPUARModelRunner)
     runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
     runner._pooler_payload_include_hidden_flag = True
@@ -283,7 +294,7 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
     runner.model = SimpleNamespace(
         prefer_model_sampler=True,
         skips_model_sampler_output_token_history=True,
-        sample=lambda logits, metadata: received.append(metadata) or "model-sampler",
+        sample=sample,
     )
     runner.sampler = SimpleNamespace()
     logits = torch.zeros((1, 4))
@@ -1670,7 +1681,6 @@ class TestMergeModelKvTransferMetadata:
         # ...while the engine-shared originals are untouched.
         assert original["r1"]["custom_metadata"] is original_meta
         assert original_meta == {"a": 1}
-        assert "talker_codes" not in original["r1"]["custom_metadata"]
 
     def test_merge_without_model_meta_passes_entry_through(self):
         class _Model:
@@ -1713,7 +1723,7 @@ class TestDownstreamPayloadMemoization:
         return runner
 
     def test_missing_marker_defaults_true_without_memoizing(self):
-        stages = {"r1": None}
+        stages: dict[str, int | None] = {"r1": None}
         runner = self._runner(stages)
 
         assert runner._request_needs_downstream_stage_payload("r1") is True
@@ -1861,6 +1871,7 @@ class TestPreferModelSamplerNoneFallback:
             "minicpmo_4_5",
             "minimax_music3",
             "nemotron_voicechat",
+            "yue2",  # Always supplies SamplerOutput, including empty/prefill steps.
         }
         assert declarers == expected, (
             "The set of models declaring `prefer_model_sampler` changed:\n"
@@ -1875,3 +1886,63 @@ class TestPreferModelSamplerNoneFallback:
             "least tolerates -- that fallback, then add its directory name to "
             "`expected` above. If you REMOVED one, drop its name."
         )
+
+
+@pytest.mark.parametrize("accepts_extra_args", [False, True])
+def test_prepare_hook_keeps_legacy_signature_and_orders_opt_in_metadata(accepts_extra_args: bool) -> None:
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = torch.nn.Module()
+    runner.model.accepts_runner_sampling_extra_args = accepts_extra_args
+    runner.requests = {
+        rid: CachedRequestState(
+            req_id=rid,
+            prompt_token_ids=[1],
+            mm_features=[],
+            sampling_params=params,
+            generator=None,
+            block_ids=([],),
+            num_computed_tokens=0,
+            output_token_ids=[],
+        )
+        for rid, params in (
+            ("a", SamplingParams(extra_args={"seed": 7})),
+            ("b", None),
+            ("unused", SamplingParams(extra_args={"seed": 9})),
+        )
+    }
+    runner.discard_request_mask = CpuGpuBuffer(3, dtype=torch.bool, device=torch.device("cpu"), pin_memory=False)
+    runner.discard_request_mask.np[:] = [True, False, True]
+    ids, positions = torch.tensor([3, 4]), torch.tensor([0, 1])
+    step_inputs = dict(
+        input_ids=ids,
+        positions=positions,
+        inputs_embeds=None,
+        num_computed_tokens=np.array([0, 1]),
+        num_scheduled_tokens=np.array([1, 1]),
+        input_ids_buffer=ids,
+    )
+    seen = {}
+
+    def legacy_hook(
+        *, input_ids, positions, inputs_embeds, req_ids, num_computed_tokens, num_scheduled_tokens, input_ids_buffer
+    ):
+        seen["req_ids"] = req_ids
+        return input_ids, positions
+
+    def opt_in_hook(*, sampling_extra_args, discard_mask, **kwargs):
+        seen["extra_args"] = sampling_extra_args
+        seen["discard"] = discard_mask.tolist()
+        return legacy_hook(**kwargs)
+
+    out = runner._call_prepare_runner_inputs(
+        opt_in_hook if accepts_extra_args else legacy_hook,
+        req_ids=["b", "a"],
+        **step_inputs,
+    )
+    assert out[0] is ids and out[1] is positions
+    assert seen["req_ids"] == ["b", "a"]
+    if accepts_extra_args:
+        assert seen["extra_args"] == [{}, {"seed": 7}]
+        assert seen["discard"] == [True, False]
+    else:
+        assert set(seen) == {"req_ids"}

tests/worker/test_omni_connector_mixin.py

@@ -39,6 +39,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class MockConnector:
+    sender_host: str | None = None
+    sender_zmq_port: int | None = None
+
     """In-memory connector for testing (mimics OmniConnectorBase)."""
 
     def __init__(self, stage_id: int = 0):
@@ -187,7 +190,7 @@ def test_init_payload_connector_ownership(role, custom_func, expected):
 
 @pytest.mark.parametrize("async_chunk", [False, True])
 @pytest.mark.parametrize("rank", [0, 1])
-def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
+def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank, mocker):
     host = MixinHost()
     host.init_omni_connectors(_make_model_config(async_chunk=async_chunk))
     # No background threads: deterministic interleaved polling.
@@ -196,7 +199,7 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
     connector = host._omni_connector
     connector.sender_host = "configured-host"
     connector.sender_zmq_port = 50051
-    connector.get = MagicMock(return_value=None)
+    get = mocker.patch.object(connector, "get", return_value=None)
     first = _make_request("r1", "external1")
     first.payload_sender_info = {"host": "producer-a", "zmq_port": "50101"}
     second = _make_request("r2", "external2")
@@ -209,14 +212,14 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
             for req_id in ("r2", "r1", "r2"):
                 assert not host._poll_single_request(req_id)
             if rank == 1:
-                connector.get.assert_not_called()
+                get.assert_not_called()
             else:
-                assert connector.get.call_args_list == [
+                assert get.call_args_list == [
                     unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}),
                     unittest.mock.call("0", "1", "external1_0_0", {"source_host": "producer-a", "source_port": 50101}),
                     unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}),
                 ]
-                connector.get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1)
+                get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1)
                 assert host._poll_single_request("r2")
                 assert "r2" not in host._pending_load_reqs
                 assert "r1" in host._pending_load_reqs
@@ -225,14 +228,14 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
             assert not host._pending_load_reqs
             # Reused internal ID without sender info cannot inherit either endpoint.
             replacement = _make_request("r1", "external-new")
-            connector.get.reset_mock(return_value=True)
-            connector.get.return_value = None
+            get.reset_mock(return_value=True)
+            get.return_value = None
             host.register_chunk_recv(replacement)
             assert not host._poll_single_request("r1")
             if rank == 0:
-                connector.get.assert_called_once_with("0", "1", "external-new_0_0")
+                get.assert_called_once_with("0", "1", "external-new_0_0")
             else:
-                connector.get.assert_not_called()
+                get.assert_not_called()
             assert connector.sender_host == "configured-host"
             assert connector.sender_zmq_port == 50051
             assert first.payload_sender_info == {"host": "producer-a", "zmq_port": "50101"}
@@ -1437,6 +1440,58 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
 
         host.shutdown_omni_connectors()
 
+    def test_first_chunk_hook_sees_only_chunk_zero_before_staging(self):
+        host = MixinHost()
+        host.init_omni_connectors(
+            model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="gen"),
+        )
+        host._omni_connector = MagicMock()
+        host._stage_id = 1
+        host._async_chunk = True
+        host._model_mode = "gen"
+        seen = []
+
+        def hook(req_id, handle, payload):
+            # The chunk is not yet visible to the scheduler when the hook runs.
+            seen.append((req_id, handle.external_req_id, payload["codes"]["audio"], "r1" in host._finished_load_reqs))
+
+        host.set_first_chunk_hook(hook)
+        host.register_chunk_recv(_make_request("r1", "ext-r1"))
+        chunk = {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(False)}}
+        host._omni_connector.get.return_value = (chunk, 1)
+        self.assertTrue(host._poll_single_request("r1"))
+        self.assertEqual(seen, [("r1", "ext-r1", [1, 2], False)])
+        self.assertIn("r1", host._finished_load_reqs)
+
+        host.get_omni_connector_output()
+        host._local_stage_payload_cache.clear()
+        host._local_request_metadata.clear()
+        host._omni_connector.get.return_value = ({"codes": {"audio": [3, 4]}, "meta": {}}, 1)
+        self.assertTrue(host._poll_single_request("r1"))
+        self.assertEqual(len(seen), 1)
+
+        host.shutdown_omni_connectors()
+
+    def test_first_chunk_hook_skips_terminal_first_chunk(self):
+        host = MixinHost()
+        host.init_omni_connectors(
+            model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="gen"),
+        )
+        host._omni_connector = MagicMock()
+        host._stage_id = 1
+        host._async_chunk = True
+        host._model_mode = "gen"
+        hook = MagicMock()
+        host.set_first_chunk_hook(hook)
+        host.register_chunk_recv(_make_request("r1"))
+        host._omni_connector.get.return_value = (
+            {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(True)}},
+            1,
+        )
+        self.assertTrue(host._poll_single_request("r1"))
+        hook.assert_not_called()
+        host.shutdown_omni_connectors()
+
     def test_non_ar_recv_waits_for_scheduler_handoff_before_fetching_next_chunk(self):
         host = MixinHost()
         host.init_omni_connectors(
@@ -1480,6 +1535,39 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
         host.shutdown_omni_connectors()
 
 
+@pytest.mark.parametrize("cancel_at", ["connector_get", "first_chunk_hook"])
+def test_cancelled_chunk_receive_does_not_recreate_delivery_state(mocker, cancel_at):
+    host = MixinHost()
+    host.init_omni_connectors(_make_model_config(async_chunk=True, worker_type="gen"))
+    host._stage_id = 1
+    host._omni_connector = mocker.MagicMock()
+    host.register_chunk_recv(_make_request("r1"))
+    seen = []
+    payload = {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(False)}}
+
+    def get(*args, **kwargs):
+        if cancel_at == "connector_get":
+            host.cleanup_finished_request("r1")
+        return payload, 1
+
+    def hook(req_id, handle, chunk):
+        seen.append(req_id)
+        host.cleanup_finished_request(req_id)
+
+    host._omni_connector.get.side_effect = get
+    host.set_first_chunk_hook(hook)
+    try:
+        assert not host._poll_single_request("r1")
+        assert seen == (["r1"] if cancel_at == "first_chunk_hook" else [])
+        assert host.get_local_stage_payload("r1") is None
+        assert "r1" not in host._get_req_chunk
+        assert "r1" not in host._pending_load_reqs
+        output = host.get_omni_connector_output()
+        assert output.chunk_ready_req_ids == set() and output.request_metadata == {}
+    finally:
+        host.shutdown_omni_connectors()
+
+
 class TestRankAwareKVRouting(unittest.TestCase):
     def _make_host(self, *, from_tp: int, to_tp: int, local_rank: int) -> MixinHost:
         host = MixinHost()

tests/worker/test_yue2_request_state.py

@@ -0,0 +1,215 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""YuE2's worker-side sampler follows retained engine outputs on recompute."""
+
+import numpy as np
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+from vllm.v1.utils import CpuGpuBuffer
+from vllm.v1.worker.gpu_input_batch import CachedRequestState
+
+from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+from vllm_omni.entrypoints.openai.tts_adapters.yue2 import Yue2Adapter
+from vllm_omni.model_executor.models.yue2 import yue2 as y
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def model(mocker):
+    """Run real state/sampling control flow with CPU staging and fixed draws."""
+    model = object.__new__(y.Yue2ForCausalLM)
+    torch.nn.Module.__init__(model)
+    model._states = {}
+    model._pending = None
+    model._step_rows = []
+    model._step_discard = None
+    model._decode_t0 = {}
+    model._last_mm = None
+    model._audio_queue = []
+    model._max_model_len = 4096
+    model._stage = torch.empty(1024, dtype=torch.long)
+    model._pinned = torch.empty((2, 64), dtype=torch.long)
+    model._synthesis = mocker.Mock(spec=y._SynthesisQueue)
+    mocker.patch.object(model, "_finish_request_safely")
+    mocker.patch.object(model, "_sampler_graph", return_value=None)
+    # CPU copies complete synchronously; this replaces only CUDA completion.
+    event = mocker.Mock(spec=torch.Event)
+    mocker.patch.object(y.torch, "Event", return_value=event)
+    mocker.patch.object(y, "sample_rows", return_value=(torch.tensor([y.CODEC_OFFSET + 7]), torch.tensor([False])))
+    constants = y._RowConstants("r", "semantic", 1.0, 0.9, 100, 1.1, 2, 1, 3, 17, False, 8)
+    model._states["r"] = y._RequestState("r", constants, prompt_len=3, prefix_ids=[1, 2, 3])
+    return model
+
+
+def test_pending_draw_is_discarded_after_chunked_recompute(model) -> None:
+    """A preempted in-flight draw never advances history or the frame limit."""
+    logits = torch.zeros((1, y.VOCAB_SIZE))
+    model._step_rows = [("r", 0, 3)]
+    model.sample(logits, None)
+    state = model._states["r"]
+    # Other requests run while r waits for KV space. Resolving the host copy
+    # may resolve a tentative token, but the resumed row must roll it back.
+    model._step_rows = []
+    model.sample(logits[:0], None)
+    assert state.job is None
+    # A partial recompute must neither accept the pending draw nor draw again.
+    model._step_rows = [("r", 0, 2)]
+    model._step_discard = [True]
+    model.sample(logits, None)
+    assert state.job is None
+    # Completing recompute has retained zero outputs: discard the old draw.
+    model._step_rows = [("r", 2, 1)]
+    model._step_discard = [False]
+    model.sample(logits, None)
+    assert state.history == []
+    # Only this new draw is retained in the next decode input.
+    model._step_rows = [("r", 3, 1)]
+    model.sample(logits, None)
+    assert state.history == [y.CODEC_OFFSET + 7]
+    assert not state.finish_ready and not state.truncated
+
+
+@pytest.mark.parametrize("token", [y.MUSIC_END, y.CODEC_OFFSET + 7])
+def test_discarded_terminal_draw_does_not_start_synthesis(model, mocker, token: int) -> None:
+    """Dropped EOS/frame-limit draws cannot prematurely finish a song."""
+    state = model._states["r"]
+    state.history = [y.CODEC_OFFSET + 2] * 2
+    mocker.patch.object(y, "sample_rows", return_value=(torch.tensor([token]), torch.tensor([False])))
+    logits = torch.zeros((1, y.VOCAB_SIZE))
+    model._step_rows = [("r", 4, 1)]
+    model.sample(logits, None)
+    model._step_rows = []
+    model.sample(logits[:0], None)
+    assert state.job is None
+    model._step_rows = [("r", 0, 5)]
+    model.sample(logits, None)
+    assert state.history == [y.CODEC_OFFSET + 2] * 2
+    assert not state.finish_ready and not state.truncated and not state.end_drawn
+    assert state.job is None and not state.finished
+    model._finish_request_safely.assert_not_called()
+
+
+def test_recompute_clears_terminal_state_and_cancels_stale_job(model) -> None:
+    """Rolling back a drawn end also releases synthesis and revives the row."""
+    state = model._states["r"]
+    state.history = [y.CODEC_OFFSET + 2] * 2
+    state.end_drawn = True
+    state.finished = state.finish_ready = True
+    state.hold_steps = 3
+    state.end_len = 6
+    state.delivered = (torch.zeros((2, 0)), False, False)
+
+    def song(keep):
+        yield
+        return torch.zeros((2, 0))
+
+    job = y._SynthesisJob("r", song)
+    state.job = job
+    model._reconcile_history(state, 2)
+    model._synthesis.cancel.assert_called_once_with(job)
+    assert state.job is None and state.hold_steps == 0
+    assert not state.finished and not state.finish_ready and not state.truncated and not state.end_drawn
+    assert state.end_len is None and state.delivered is None
+
+
+def test_dropped_end_step_redelivers_song_and_end(model, mocker) -> None:
+    """A preemption that drops the step emitting the end token also drops
+    that step's audio; the resumed request must deliver both again."""
+    state = model._states["r"]
+    state.history = [y.CODEC_OFFSET + 2] * 3  # at the 3-frame budget
+    state.truncated = True
+    state.hold_steps = 2
+    audio = torch.ones((2, 5))
+    job = mocker.Mock(t_start=1.0, t_done=2.0)
+    job.done.return_value = True
+    state.job = job
+    model._synthesis.complete.return_value = (audio, False, False)
+    logits = torch.zeros((1, y.VOCAB_SIZE))
+
+    # Engine holds 3 frames + 2 HOLDs; this step ships the song and its end.
+    model._step_rows = [("r", 7, 1)]
+    out = model.sample(logits, None)
+    assert out.sampled_token_ids.tolist() == [[y.MUSIC_END]]
+    assert state.finished and state.end_len == 6
+    assert [entry[:2] for entry in model._audio_queue] == [("r", audio)]
+
+    # An async lookahead row (end token in flight) delivers nothing again.
+    model._step_rows = [("r", 8, 1)]
+    model.sample(logits, None)
+    assert len(model._audio_queue) == 1
+
+    # Preempted with the end step in flight: the resumed row lacks it.
+    model._step_rows = [("r", 0, 8)]
+    model._step_discard = [False]
+    out = model.sample(logits, None)
+    assert out.sampled_token_ids.tolist() == [[y.MUSIC_END]]
+    assert [entry[:2] for entry in model._audio_queue] == [("r", audio), ("r", audio)]
+    assert model._audio_queue[1][2:] == (True, False)  # truncated, not failed
+    assert state.end_len == 6
+    model._synthesis.complete.assert_called_once()
+
+
+def test_accepted_draw_still_reaches_frame_limit(model) -> None:
+    """Confirmed codec frames retain the existing delayed finish behavior."""
+    logits = torch.zeros((1, y.VOCAB_SIZE))
+    for comp, span in ((0, 3), (3, 1), (4, 1)):
+        model._step_rows = [("r", comp, span)]
+        model.sample(logits, None)
+    state = model._states["r"]
+    assert len(state.history) == 2
+    model._resolve_pending()
+    model._reconcile_history(state, 3)
+    assert len(state.history) == 3
+    assert state.finish_ready and state.truncated
+
+
+def test_real_prepare_hook_admits_reordered_requests(model) -> None:
+    """Runner opt-in metadata reaches the actual YuE2 admission hook."""
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = model
+    runner.requests = {}
+    for rid, seed in (("a", 7), ("b", 9)):
+        runner.requests[rid] = CachedRequestState(
+            req_id=rid,
+            prompt_token_ids=[1, 2, 3],
+            mm_features=[],
+            sampling_params=SamplingParams(extra_args={y.KEY_PREFIX_IDS: [1, 2, 3], y.KEY_SEED: seed}),
+            generator=None,
+            block_ids=([],),
+            num_computed_tokens=0,
+            output_token_ids=[],
+        )
+    runner.discard_request_mask = CpuGpuBuffer(2, dtype=torch.bool, device=torch.device("cpu"), pin_memory=False)
+    runner.discard_request_mask.np[:] = [True, False]
+    ids = torch.tensor([1, 2, 1, 2, 3], dtype=torch.int32)
+    positions = torch.tensor([0, 1, 0, 1, 2])
+    runner._call_prepare_runner_inputs(
+        model.prepare_runner_inputs,
+        req_ids=["b", "a"],
+        input_ids=ids,
+        positions=positions,
+        inputs_embeds=None,
+        num_computed_tokens=np.array([0, 0]),
+        num_scheduled_tokens=np.array([2, 3]),
+        input_ids_buffer=ids,
+    )
+    assert model._states["b"].constants.seed == 9
+    assert model._states["a"].constants.seed == 7
+    assert model._states["b"].prompt_len == 3
+    assert model._states["b"].prefix_ids == [1, 2, 3]
+    assert model._step_discard == [True, False]
+
+
+@pytest.mark.parametrize("spare", [0, 1, 4095, 4096, 5000])
+def test_hold_budget_matches_engine_headroom(spare: int) -> None:
+    """Adapter metadata advertises only the HOLD steps left after clipping."""
+    frames = 250
+    prompt = {"prompt_token_ids": [1] * (y.CONTEXT - frames - 1 - spare)}
+    adapter = object.__new__(Yue2Adapter)
+    request = OpenAICreateSpeechRequest(model="m-a-p/YuE2-3B", input="lyrics", max_new_tokens=frames)
+    params = adapter.apply_sampling_overrides([SamplingParams()], request, prompt)[0]
+    assert params.extra_args[y.KEY_MAX_HOLD_STEPS] == min(spare, 4096)
+    assert params.max_tokens == frames + 1 + min(spare, 4096)

tests/worker/test_yue2_sampler.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""YuE2's batched sampler draws exactly what the per-row reference draws.
+
+The model samples all rows of a preset as one batch, over the phase's own
+columns only, and writes the multinomial out as ``argmax(p / q)`` with
+``q ~ Exp(1)`` so the draw stays on the device. The reference below is the
+upstream per-row arithmetic over the full vocabulary with
+``torch.multinomial``. Same seeds must give the same tokens, and each
+request's generator must advance by exactly the same amount per step, or a
+request's song changes with the batch it happens to share.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.yue2.yue2 import (
+    ABC_END,
+    ABC_SAMPLING,
+    CODEC_OFFSET,
+    CODEC_SIZE,
+    EOD,
+    MUSIC_END,
+    PHASE_VOCAB,
+    SEMANTIC_SAMPLING,
+    VOCAB_SIZE,
+    sample_rows,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _reference_scores(logits: torch.Tensor, preset: dict, history: list[int], phase: str) -> torch.Tensor:
+    """Upstream per-row scores over the full vocabulary (float32)."""
+    scores = logits.float()  # a new tensor: the logits are bf16
+    end = ABC_END if phase == "abc" else MUSIC_END
+    allowed = torch.full_like(scores, -torch.inf)
+    if phase == "abc":
+        allowed[:EOD] = 0
+    else:
+        allowed[CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0
+    allowed[end] = 0
+    scores = scores + allowed
+    if len(history) < preset["min_tokens"]:
+        scores[end] = -torch.inf
+    recent = history[-preset["penalty_window"] :]
+    if preset["repetition_penalty"] != 1.0 and recent:
+        freq = torch.zeros_like(scores)
+        freq.scatter_add_(-1, torch.tensor(recent), torch.ones(len(recent)))
+        alpha = preset["repetition_penalty"] ** freq
+        scores = torch.where(scores < 0, scores * alpha, scores / alpha)
+    if preset["temperature"] == 0:
+        return scores
+    scores = scores / preset["temperature"]
+    threshold = scores.topk(preset["top_k"]).values[-1]
+    scores = scores.masked_fill(scores < threshold, -torch.inf)
+    values, indices = scores.sort(descending=True, stable=True)
+    probabilities = values.softmax(-1)
+    removed = probabilities.cumsum(-1) - probabilities > preset["top_p"]
+    removed[:1] = False
+    values = values.masked_fill(removed, -torch.inf)
+    return values.scatter(-1, indices, values)
+
+
+def _reference_draw(scores: torch.Tensor, generator: torch.Generator, greedy: bool) -> int:
+    if greedy:
+        return int(scores.argmax())
+    return int(torch.multinomial(scores.softmax(-1), 1, generator=generator))
+
+
+def _batched_draw(logits, preset, histories, phase, generators) -> list[int]:
+    vocab = PHASE_VOCAB[phase]
+    window = preset["penalty_window"]
+    recent = [[vocab.col(t) for t in history[-window:]] for history in histories]
+    width = max(map(len, recent))
+    window_cols = torch.tensor([r + [vocab.num_cols] * (width - len(r)) for r in recent]) if width else None
+    block = [len(history) < preset["min_tokens"] for history in histories]
+    ids, bad = sample_rows(
+        logits,
+        vocab,
+        temperature=preset["temperature"],
+        top_p=preset["top_p"],
+        top_k=preset["top_k"],
+        repetition_penalty=preset["repetition_penalty"],
+        window_cols=window_cols,
+        block_end=torch.tensor(block) if any(block) else None,
+        generators=generators,
+    )
+    assert not bad.any()
+    return ids.tolist()
+
+
+@pytest.mark.parametrize(
+    ("phase", "rows", "min_tokens", "temperature"),
+    [
+        ("semantic", 1, None, None),
+        ("semantic", 3, None, None),
+        ("semantic", 3, 0, None),  # the end token may be drawn
+        ("semantic", 2, None, 0.0),  # greedy
+        ("abc", 2, None, None),
+        ("abc", 2, 0, None),
+    ],
+)
+def test_batched_draws_match_per_row_multinomial(phase, rows, min_tokens, temperature) -> None:
+    preset = dict(ABC_SAMPLING if phase == "abc" else SEMANTIC_SAMPLING)
+    if min_tokens is not None:
+        preset["min_tokens"] = min_tokens
+    if temperature is not None:
+        preset["temperature"] = temperature
+    end = ABC_END if phase == "abc" else MUSIC_END
+    reference = [torch.Generator().manual_seed(1000 + row) for row in range(rows)]
+    batched = [torch.Generator().manual_seed(1000 + row) for row in range(rows)]
+    histories: list[list[int]] = [[] for _ in range(rows)]
+    logits_generator = torch.Generator().manual_seed(7)
+    for _ in range(40):
+        logits = (torch.randn((rows, VOCAB_SIZE), generator=logits_generator) * 3).to(torch.bfloat16)
+        # Peaky over a few ids, like a real LM, so the penalty and top-p matter.
+        if phase == "semantic":
+            logits[:, CODEC_OFFSET : CODEC_OFFSET + 300] += 6
+        else:
+            logits[:, :300] += 6
+        logits[:, end] += 4
+        expected = [
+            _reference_draw(
+                _reference_scores(logits[row], preset, histories[row], phase), reference[row], temperature == 0
+            )
+            for row in range(rows)
+        ]
+        assert _batched_draw(logits, preset, histories, phase, batched) == expected
+        for row in range(rows):
+            # A greedy request never reads its stream (its temperature is
+            # fixed for life), so only sampling rows must advance in step.
+            if temperature != 0:
+                assert torch.equal(reference[row].get_state(), batched[row].get_state())
+            if expected[row] != end:
+                histories[row].append(expected[row])
+
+
+def test_a_row_draws_the_same_alone_or_in_a_batch() -> None:
+    """A request's song must not depend on which requests share its step."""
+    preset = SEMANTIC_SAMPLING
+    logits = torch.randn((3, VOCAB_SIZE), generator=torch.Generator().manual_seed(3)).to(torch.bfloat16)
+    histories = [[CODEC_OFFSET + i for i in range(60)], [], [CODEC_OFFSET + 5] * 10]
+    together = _batched_draw(
+        logits, preset, histories, "semantic", [torch.Generator().manual_seed(s) for s in (1, 2, 3)]
+    )
+    alone = [
+        _batched_draw(
+            logits[row : row + 1], preset, histories[row : row + 1], "semantic", [torch.Generator().manual_seed(s)]
+        )[0]
+        for row, s in enumerate((1, 2, 3))
+    ]
+    assert together == alone
+
+
+@pytest.mark.parametrize("phase", ["abc", "semantic"])
+def test_phase_vocab_columns_round_trip(phase) -> None:
+    vocab = PHASE_VOCAB[phase]
+    end = ABC_END if phase == "abc" else MUSIC_END
+    ids = [end, 0, EOD - 1] if phase == "abc" else [end, CODEC_OFFSET, CODEC_OFFSET + CODEC_SIZE - 1]
+    cols = torch.tensor([vocab.col(i) for i in ids])
+    assert vocab.col(end) == vocab.end_col
+    assert cols.max() < vocab.num_cols
+    assert vocab.ids(cols).tolist() == ids
+    logits = torch.randn((2, VOCAB_SIZE))
+    assert torch.equal(vocab.take(logits)[:, cols], logits[:, ids])