[Model] YuE2: async single-stage synthesis and CUDA graph acceleration (#8407)

Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: Wang Zhipeng <wangzhipeng628@gmail.com>
Co-authored-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

examples/offline_inference/yue2/end2end.py

@@ -28,7 +28,9 @@ import torch
 from vllm_omni import Omni
 from vllm_omni.model_executor.models.yue2.yue2 import (
     ABC_SAMPLING,
+    CONTEXT,
     KEY_MAX_AUDIO_FRAMES,
+    KEY_MAX_HOLD_STEPS,
     KEY_MIN_TOKENS,
     KEY_PENALTY_WINDOW,
     KEY_PHASE,
@@ -41,10 +43,12 @@ from vllm_omni.model_executor.models.yue2.yue2 import (
     KEY_TOP_P,
     SEMANTIC_SAMPLING,
     STOP_TOKEN_IDS,
+    SYNTHESIS_HOLD_STEPS,
 )
 from vllm_omni.tokenizers.yue2_prompt import (
     abc_ids_from_generated,
     abc_prefix_ids,
+    semantic_frames,
     semantic_prefix_ids,
 )
 from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer
@@ -71,7 +75,13 @@ def sampling_params(engine, *, phase, seed, max_frames, prompt_ids):
         # conditioning prefix from the scheduled tokens alone.
         KEY_PREFIX_IDS: list(prompt_ids),
     }
-    params.max_tokens = preset["max_tokens"] if phase == "abc" else max_frames + 1
+    if phase == "abc":
+        params.max_tokens = preset["max_tokens"]
+    else:
+        # The song is synthesized on a side stream while the request is held;
+        # leave max_tokens headroom for the hold tokens.
+        params.max_tokens = min(max_frames + 1 + SYNTHESIS_HOLD_STEPS, CONTEXT - len(prompt_ids))
+        params.extra_args[KEY_MAX_HOLD_STEPS] = max(0, params.max_tokens - max_frames - 1)
     params.stop_token_ids = list(STOP_TOKEN_IDS)
     params.detokenize = False
     return params
@@ -144,6 +154,9 @@ def main() -> None:
         payload = {
             "prompt_token_ids": prompt_ids,
             "generated_token_ids": generated_ids,
+            # Codec frames only: without the end token and the hold tokens
+            # the request emits while its song is synthesized.
+            "semantic_frames": semantic_frames(generated_ids),
             "seed": args.seed,
             "cot": args.cot,
             "max_frames": args.max_frames,

recipes/m-a-p/YuE2-3B.md

@@ -40,9 +40,11 @@ user-downloaded weights and bundles none.
 | Output | 48 kHz stereo, whole song delivered when the semantic request finishes |
 
 Generation is one or two requests on a single AR stage: the abc phase writes
-the score, the semantic phase generates codec frames and, on its last step,
-solves the ODE and decodes the whole song in-engine (weights load once; the
-VAE loads at startup from `$YUE2_VAE` or the hub id).
+the score, the semantic phase generates codec frames, then queues the ODE and VAE on a
+high-priority CUDA stream. Other requests keep decoding while the finishing
+request emits HOLD tokens (up to 4096 steps, within its token/context budget).
+The scheduler is unchanged. Weights load once; the VAE loads at startup from
+`$YUE2_VAE` or the hub id.
 
 ## Running
 
@@ -67,15 +69,50 @@ ending.
 
 ## Notes
 
-- **Memory (RTX 4090, 24 GB, `gpu_memory_utilization: 0.70`):** the engine
-  reserves ~15.8 GiB after startup; with 4 concurrent requests pinned to the
-  9000-frame cap (worst case), peak usage during the terminal NAR/VAE
-  finishing pass reaches ~20.9 GiB, leaving ~3 GiB of headroom. At 0.85 the
-  same workload OOMs inside the finishing pass, so the deploy yaml pins 0.70.
-- **Known limitations:** 4 concurrent full-length requests are the verified
-  shape (`max_num_seqs: 4`); CUDA graph capture is a
-  follow-up; the abc phase runs eagerly
-  after prefill (its tokens are the product, not audio).
+The lifecycle e2e regression forces KV-cache preemption with chunked and
+unchunked recompute, mixes ABC and semantic sampling, and cancels a request
+after NAR work has started. It checks history alignment, released synthesis
+buffers, and successful generation after cancellation. Run with one H100/H200
+and cached YuE2-3B and YuE2-Vae weights:
+
+```bash
+export YUE2_MODEL_DIR=/path/to/YuE2-3B
+export YUE2_VAE=/path/to/YuE2-Vae
+CUDA_VISIBLE_DEVICES=0 python -m pytest tests/e2e/online_serving/test_yue2.py \
+    -k preemption_and_synthesis_abort -m 'slow and tts' --run-level full_model -q -s
+```
+
+The test retains the module's weekly TTS routing; it is not a per-PR CI gate.
+
+- **Memory:** `gpu_memory_utilization` budgets the vLLM engine; NAR K/V,
+  acoustic graph buffers and VAE activations also need room during synthesis.
+  The 0.70 / 4-slot defaults require validation on the target card. The earlier
+  RTX 4090 measurement (15.8 GiB after startup, 20.9 GiB peak at 9000 frames)
+  predates the async/compiled implementation; it does not validate this
+  version or long ABC prefixes on a 4090. Reduce the fraction if the target
+  workload needs more synthesis memory, especially for long ABC prefixes and
+  the 9000-frame cap. Releasing completed chunks bounds live NAR buffers,
+  while the shared graph allocator retains reserved storage for reuse.
+- **CUDA graphs:** the default deploy keeps the AR backbone eager. Both ABC
+  and semantic sampling use prewarmed power-of-two graph buckets. Custom
+  sampler captures have a bounded cache and fall back to eager sampling when
+  it fills. The NAR velocity pass uses a chunk graph on the serialized synthesis
+  stream. Chunk graphs share one allocator pool; completed chunks release
+  their engines and buffers after their own events. FA3 runs on Hopper; other
+  CUDA cards use SDPA, which is also warmed at startup.
+- **Serving:** run `vllm serve m-a-p/YuE2-3B --omni` with the default
+  `yue2.yaml`, which uses eager AR with 4 slots. Validate memory capacity
+  for the target card and workload.
+  For the tested single-H200 graph settings, reuse this deploy with
+  `--stage-overrides '{"0":{"max_num_seqs":32,"gpu_memory_utilization":0.5,"enforce_eager":false,"compilation_config":{"cudagraph_mode":"FULL_AND_PIECEWISE"}}}'`.
+  Decode uses FULL and mixed prefill uses PIECEWISE; VAE and the Python
+  synthesis queue remain outside these graphs.
+  Audio is delivered as a whole song; time to first audio equals completion
+  latency. Aborting a running song stops future work units and waits only for
+  the at-most-two submitted units before releasing its buffers.
+- **Preemption:** before drawing for a resumed row, model-owned sampling
+  reconciles the codec history and terminal flags with the engine's retained
+  sequence. Discarded draws cannot advance its penalty window or frame budget.
 - Audio is not expected to match the upstream torch reference bit-for-bit
   (fused vs eager kernels). Measured against the upstream torch reference
   (bd90e4c, same style/lyrics/seed 831001, cot=off, 200 frames): prompt token

tests/e2e/features/yue2/worker_extension.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Observe YuE2 lifecycle contracts inside a real engine worker."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+
+class Yue2LifecycleWorkerExtension:
+    model_runner: "GPUARModelRunner"
+
+    def start_yue2_probe(self):
+        from unittest.mock import patch
+
+        from vllm.compilation.cuda_graph import CUDAGraphWrapper
+        from vllm.config import CUDAGraphMode
+        from vllm.forward_context import get_forward_context, is_forward_context_available
+
+        from vllm_omni.model_executor.models.yue2.yue2 import ABC_END, HOLD_TOKEN, MUSIC_END
+
+        model = self.model_runner.get_model()
+        self.yue2_probe = {
+            "preemptions": 0,
+            "rollbacks": 0,
+            "checked_rows": 0,
+            "mixed_steps": 0,
+            "synthesis_started": False,
+            "ar_full_replays": 0,
+            "ar_piecewise_replays": 0,
+        }
+        graph_call = CUDAGraphWrapper.__call__
+
+        def count_ar_replays(wrapper, *args, **kwargs):
+            if is_forward_context_available():
+                context = get_forward_context()
+                mode = context.cudagraph_runtime_mode
+                entry = wrapper.concrete_cudagraph_entries.get(context.batch_descriptor)
+                if mode == wrapper.runtime_mode and entry is not None and entry.cudagraph is not None:
+                    if mode == CUDAGraphMode.FULL:
+                        self.yue2_probe["ar_full_replays"] += 1
+                    elif mode == CUDAGraphMode.PIECEWISE:
+                        self.yue2_probe["ar_piecewise_replays"] += 1
+            return graph_call(wrapper, *args, **kwargs)
+
+        # This worker is created for this test and exits with the engine.
+        # Count vLLM AR replays, separately from model-owned sampler/NAR graphs.
+        self.yue2_graph_patch = patch.object(CUDAGraphWrapper, "__call__", count_ar_replays)
+        self.yue2_graph_patch.start()
+        self.yue2_cancel_job = None
+        previous: dict[str, int] = {}
+        prepare = model.prepare_runner_inputs
+        reconcile = model._reconcile_history
+        queue = model._queue_synthesis
+        sample = model.sample
+        resolve = model._resolve_pending
+        emitted: dict[str, int] = {}
+
+        def resolve_pending():
+            if model._pending is not None:
+                event, host, states = model._pending
+                event.synchronize()
+                for state, token in zip(states, host[0].tolist()):
+                    end = ABC_END if state.constants.phase == "abc" else MUSIC_END
+                    expected = HOLD_TOKEN if token == end and not state.constants.skip_synthesis else token
+                    assert emitted[state.request_id] == expected, "engine token and model history disagree"
+            resolve()
+
+        def sample_tokens(logits, sampling_metadata):
+            result = sample(logits, sampling_metadata)
+            phases = {model._states[rid].constants.phase for rid, _, _ in model._step_rows if rid in model._states}
+            if len(phases) == 2:
+                self.yue2_probe["mixed_steps"] += 1
+            for (rid, _, _), row in zip(model._step_rows, result.sampled_token_ids.tolist()):
+                emitted[rid] = row[0]
+            return result
+
+        def prepare_inputs(**kwargs):
+            for rid, computed in zip(kwargs["req_ids"], kwargs["num_computed_tokens"]):
+                if int(computed) < previous.get(rid, 0):
+                    self.yue2_probe["preemptions"] += 1
+                previous[rid] = int(computed)
+            return prepare(**kwargs)
+
+        def reconcile_history(state, accepted):
+            if len(state.history) + int(state.end_drawn) > accepted:
+                self.yue2_probe["rollbacks"] += 1
+            reconcile(state, accepted)
+            assert len(state.history) + int(state.end_drawn) <= accepted
+            if not state.finish_ready and state.job is None and not state.finished:
+                assert len(state.history) == accepted
+                assert not state.truncated and not state.end_drawn
+            self.yue2_probe["checked_rows"] += 1
+
+        def queue_synthesis(state):
+            queue(state)
+            if "cancel-synthesis" in state.request_id:
+                job = state.job
+                assert job.started and job.work is not None and job.keep
+                assert job.event is None, "test must intercept an active synthesis job"
+                # Keep an actual NAR job active until the driver cancels it.
+                # Its first work units ran normally; stop pumping later units
+                # so cancellation cannot race against a short song finishing.
+                self.yue2_cancel_advance = job.advance
+                job.advance = lambda: None
+                self.yue2_cancel_job = job
+                self.yue2_probe["synthesis_started"] = True
+
+        model.prepare_runner_inputs = prepare_inputs
+        model._reconcile_history = reconcile_history
+        model._queue_synthesis = queue_synthesis
+        model.sample = sample_tokens
+        model._resolve_pending = resolve_pending
+        return True
+
+    def get_yue2_probe(self):
+        result = dict(self.yue2_probe)
+        job = self.yue2_cancel_job
+        if job is not None:
+            if job.cancelled:
+                # Exercise the original pump after cleanup: it must not submit
+                # further units even though this test retains the job object.
+                self.yue2_cancel_advance()
+            model = self.model_runner.get_model()
+            result["cancelled"] = job.cancelled
+            result["released"] = job._released and job.work is None and not job.keep and job.stream is None
+            result["removed"] = (
+                job.request_id not in model._states
+                and model._synthesis.active is not job
+                and job not in model._synthesis.waiting
+            )
+        return result

tests/e2e/online_serving/test_yue2.py

@@ -32,9 +32,9 @@ above the byte floor) and passes on that alone, following the MiniMax Music 3
 precedent.
 
 Weekly rather than per-PR: each request is a 16 s song and the verification
-checkpoints add ~5.5 GB of downloads. Validated on a single RTX 4090
-(24 GB); the hardware mark targets an H100 because the CI fleet has no
-consumer cards.
+checkpoints add ~5.5 GB of downloads. The async/graph implementation was
+validated on H200; the hardware mark targets the CI fleet's H100. The older
+RTX 4090 result does not establish this implementation's memory requirements.
 """
 
 import concurrent.futures
@@ -60,8 +60,7 @@ DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 900.0
 # 400 frames = 16 s of song. The byte floor sits far below the observed
 # payload (~3.0 MB of 48 kHz stereo 16-bit) to leave room for an early EOS.
 _FRAMES = 400
-# Concurrent pair stays short: it exists to exercise row re-indexing, not
-# to double the wall time.
+# Short concurrent requests exercise row re-indexing as budgets differ.
 _CONCURRENT_FRAMES = 250
 _MIN_BYTES = 1_000_000
 
@@ -173,7 +172,7 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
     # CI runs without the verification assets, so the structural checks below
     # are the only guard against the two real bugs this PR fixed during review:
     # mono/half-speed output (WAV header + duration bounds) and the int32
-    # crash under concurrent row re-indexing (two simultaneous requests).
+    # crash under concurrent row re-indexing (simultaneous requests).
     # Duration bounds: the end token is masked for the first 200 steps
     # (min_tokens), and the 400-frame budget caps the song, so a healthy
     # request lands in [8, 16.5] s; the half-speed bug produced ~32 s.
@@ -190,13 +189,15 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
     # (see _http_speech); the OpenAI SDK response does not expose headers.
     _http_speech(omni_server, seed=SEED, frames=_FRAMES)
 
-    # Two concurrent requests with different seeds: both must succeed and
-    # produce different songs (the model-owned sampler re-indexes rows across
-    # concurrent requests; a dtype mismatch here crashed the engine before).
-    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
-        futures = [pool.submit(_http_speech, omni_server, seed=SEED + i, frames=_CONCURRENT_FRAMES) for i in range(2)]
-        pair = [f.result() for f in futures]
-    assert pair[0] != pair[1], "different seeds produced identical audio"
+    # Different budgets exercise row re-indexing as shorter requests finish.
+    # Compare the same-budget pair so length cannot hide an ignored seed.
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as pool:
+        futures = [
+            pool.submit(_http_speech, omni_server, seed=SEED + i, frames=frames)
+            for i, frames in enumerate((_CONCURRENT_FRAMES, _FRAMES, _CONCURRENT_FRAMES))
+        ]
+        audios = [f.result() for f in futures]
+    assert audios[0] != audios[2], "different seeds produced identical audio at the same frame budget"
 
     # Without the verification assets the structural assertions above
     # (non-empty payload, byte floor) are the whole test — pass on them.
@@ -235,3 +236,199 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
     assert melody_score > control_score, (
         f"calibration: ABC-conditioned {melody_score:.3f} not above control {control_score:.3f}"
     )
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_yue2_abort_then_next_request_succeeds(omni_server) -> None:
+    """Serving smoke check after a client timeout, not proof of job cleanup."""
+    # The chosen seed and budget time out on the tested cards. Require the
+    # timeout so this cannot silently become another successful request;
+    # whether it lands in AR or synthesis depends on the card.
+    with pytest.raises(httpx.TimeoutException):
+        httpx.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech",
+            json={
+                "model": omni_server.model,
+                "input": LYRICS,
+                "instructions": CAPTION,
+                "seed": SEED + 2,
+                "max_new_tokens": 3000,
+                "stream": False,
+                "response_format": "wav",
+                "extra_params": {"cot": "off"},
+            },
+            timeout=1.0,
+        )
+    _http_speech(omni_server, seed=SEED + 3, frames=_CONCURRENT_FRAMES)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("chunked", [True, False], ids=["chunked", "unchunked"])
+@pytest.mark.parametrize("enforce_eager", [True, False], ids=["eager", "ar-graphs"])
+async def test_yue2_preemption_and_synthesis_abort(chunked, enforce_eager, monkeypatch) -> None:
+    """Force real preemption and abort only after NAR work has been submitted."""
+    import asyncio
+    from pathlib import Path
+
+    import torch
+    from vllm import SamplingParams
+
+    from tests.helpers.runtime import get_model_prefix
+    from tests.helpers.stage_config import modify_stage_config
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.model_executor.models.yue2.yue2 import (
+        ABC_END,
+        EOD,
+        KEY_MAX_AUDIO_FRAMES,
+        KEY_MAX_HOLD_STEPS,
+        KEY_MIN_TOKENS,
+        KEY_PHASE,
+        KEY_PREFIX_IDS,
+        KEY_SEED,
+        KEY_SKIP_SYNTHESIS,
+        STOP_TOKEN_IDS,
+    )
+    from vllm_omni.tokenizers.yue2_prompt import abc_prefix_ids, semantic_frames, semantic_prefix_ids
+    from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer
+    from vllm_omni.transformers_utils.repo_utils import hf_api
+
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    model_path = Path(os.environ.get("YUE2_MODEL_DIR", get_model_prefix() + MODEL))
+    if not model_path.is_dir():
+        model_path = Path(hf_api().snapshot_download(str(model_path)))
+    tokenizer = YuE2TextTokenizer(model_path / "qwen.tiktoken")
+    prefix = semantic_prefix_ids(tokenizer.encode, CAPTION, LYRICS, "off")
+    frames = 256
+    # Each song fits alone, but four full songs cannot fit in 64 KV blocks.
+    # Chunked recompute also replays output tokens past the original prompt.
+    deploy = modify_stage_config(
+        get_deploy_config_path("yue2.yaml"),
+        updates={
+            "stages": {
+                0: {
+                    "enforce_eager": enforce_eager,
+                    "compilation_config": {
+                        "cudagraph_mode": "NONE" if enforce_eager else "FULL_AND_PIECEWISE",
+                    },
+                    "enable_chunked_prefill": chunked,
+                    "async_scheduling": True,
+                    "max_num_batched_tokens": 256 if chunked else 1024,
+                    "max_model_len": 1000,
+                    "num_gpu_blocks_override": 64,
+                    "block_size": 16,
+                    "max_num_seqs": 4,
+                    "tokenizer": str(model_path),
+                }
+            }
+        },
+    )
+    engine = AsyncOmni(
+        model=str(model_path),
+        deploy_config=deploy,
+        worker_extension_cls="tests.e2e.features.yue2.worker_extension.Yue2LifecycleWorkerExtension",
+        stage_init_timeout=600,
+        init_timeout=900,
+    )
+
+    async def probe():
+        return (await engine.collective_rpc("get_yue2_probe", stage_ids=[0]))[0][0]
+
+    async def song(request_id, seed):
+        params = SamplingParams(
+            max_tokens=1000 - len(prefix),
+            stop_token_ids=STOP_TOKEN_IDS,
+            detokenize=False,
+            extra_args={
+                KEY_PREFIX_IDS: prefix,
+                KEY_SEED: seed,
+                KEY_MAX_AUDIO_FRAMES: frames,
+                KEY_MIN_TOKENS: frames,
+                KEY_MAX_HOLD_STEPS: 1000 - len(prefix) - frames - 1,
+            },
+        )
+        audio = []
+        tokens = []
+        async for output in engine.generate(
+            {"prompt_token_ids": prefix}, request_id=request_id, sampling_params_list=[params]
+        ):
+            tokens = list(output.outputs[0].token_ids)
+            multimodal = getattr(output.outputs[0], "multimodal_output", None)
+            if multimodal is not None and "audio" in multimodal:
+                data = multimodal["audio"]
+                audio.extend(data if isinstance(data, list) else [data])
+        assert audio, "request completed without audio"
+        waveform = torch.cat([part.detach().cpu().reshape(-1) for part in audio])
+        assert torch.isfinite(waveform).all() and waveform.abs().max() > 0
+        assert len(semantic_frames(tokens)) == frames
+        # The VAE convolution trims a small number of boundary samples.
+        assert waveform.numel() == pytest.approx(frames * 1920 * 2, abs=256)
+
+    async def plan():
+        abc_prefix = abc_prefix_ids(tokenizer.encode, CAPTION, LYRICS, "melody")
+        params = SamplingParams(
+            max_tokens=64,
+            stop_token_ids=STOP_TOKEN_IDS,
+            detokenize=False,
+            extra_args={
+                KEY_PHASE: "abc",
+                KEY_PREFIX_IDS: abc_prefix,
+                KEY_SEED: SEED,
+                KEY_SKIP_SYNTHESIS: True,
+            },
+        )
+        tokens = []
+        async for output in engine.generate(
+            {"prompt_token_ids": abc_prefix}, request_id="abc-plan", sampling_params_list=[params]
+        ):
+            tokens = list(output.outputs[0].token_ids)
+        assert tokens and all(t < EOD or t == ABC_END for t in tokens)
+
+    tasks = []
+    try:
+        assert await engine.collective_rpc("start_yue2_probe", stage_ids=[0]) == [[True]]
+        tasks = [asyncio.create_task(plan())]
+        tasks.extend(asyncio.create_task(song(f"preempt-{i}", SEED + i)) for i in range(4))
+        await asyncio.wait_for(asyncio.gather(*tasks), timeout=600)
+        observations = await probe()
+        assert observations["preemptions"] > 0, observations
+        assert observations["rollbacks"] > 0, observations
+        assert observations["checked_rows"] >= 4 * frames, observations
+        assert observations["mixed_steps"] > 0, observations
+        if enforce_eager:
+            assert observations["ar_full_replays"] == 0, observations
+        else:
+            assert observations["ar_full_replays"] > 0, observations
+
+        cancelled = asyncio.create_task(song("cancel-synthesis", SEED + 7))
+        tasks.append(cancelled)
+        deadline = asyncio.get_running_loop().time() + 120
+        while not (await probe())["synthesis_started"]:
+            assert not cancelled.done(), "song finished before the active-job cancellation probe"
+            assert asyncio.get_running_loop().time() < deadline, "NAR job never started"
+            await asyncio.sleep(0.05)
+        cancelled.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await cancelled
+        deadline = asyncio.get_running_loop().time() + 30
+        while True:
+            observations = await probe()
+            if observations.get("removed"):
+                assert observations["cancelled"] and observations["released"], observations
+                break
+            assert asyncio.get_running_loop().time() < deadline, observations
+            await asyncio.sleep(0.05)
+        await asyncio.wait_for(song("after-cancel", SEED + 8), timeout=120)
+        observations = await probe()
+        assert observations["released"]
+        print("YuE2 lifecycle:", observations)
+    finally:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        engine.shutdown()

tests/worker/test_gpu_ar_model_runner.py

@@ -16,6 +16,7 @@ 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
 
@@ -1885,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_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])

vllm_omni/deploy/yue2.yaml

@@ -1,5 +1,5 @@
 # YuE2-3B deploy: single-stage native-AR text-to-music (~3.6B bf16 + VAE).
-# Target 1x 24 GB card; verified on a single RTX 4090.
+# Four-slot CUDA configuration; validate memory on the target card/workload.
 #
 # The architectures pin is required: the checkpoint's config.json names
 # "YuE2ForCausalLM" with an auto_map that only makes sense under
@@ -13,16 +13,17 @@ dtype: bfloat16
 
 stages:
   - stage_id: 0
+    # A finishing song keeps its slot during side-stream synthesis,
+    # emitting HOLD tokens until its audio is ready.
     max_num_seqs: 4
-    # Measured on a 24 GB RTX 4090 with 4 concurrent requests pinned to the
-    # 9000-frame cap (worst case): the engine reserves ~15.8 GiB at 0.70 and
-    # peak usage during the terminal NAR/VAE finishing pass reaches 20.9 GiB,
-    # leaving ~3 GiB of headroom. 0.85 reserves ~19.4 GiB and the same
-    # workload OOMs inside the finishing pass.
+    # Reserve synthesis headroom: NAR K/V, acoustic graph buffers and VAE
+    # activations are allocated beyond the engine budget. The earlier 4090
+    # memory measurement predates async/compiled synthesis; 0.70 is a starting
+    # point, not a guarantee for 9000 frames or long ABC prefixes. Lower it if
+    # the target workload needs more synthesis memory (see the model recipe).
     gpu_memory_utilization: 0.70
-    # Phase 1 keeps the backbone eager: the model-owned sampler runs after the
-    # forward, and the terminal NAR/VAE pass runs inside the finishing step.
-    # CUDA-graph capture of the standard Qwen3 backbone is a perf follow-up.
+    # Keep AR capture off by default. Sampler/NAR graphs still need
+    # separate synthesis headroom.
     enforce_eager: true
     # Prefix caching stays off: the omni prefix-cache save snapshots the
     # step's mm payload before the model-owned sampler runs, but YuE2
@@ -42,14 +43,17 @@ stages:
     hf_overrides:
       architectures: ["Yue2ForCausalLM"]
     default_sampling_params:
-      # Placeholder values: the model's own sampler overrides everything from
-      # extra_args. Their presence turns on has_sampling_extra_args, which is
-      # what threads the per-request yue2_* keys into the model's forward.
+      # Defaults for callers that omit model-owned sampling options. The
+      # opt-in prepare_runner_inputs hook reads each request's extra_args;
+      # it does not depend on the forward's has_sampling_extra_args flag.
       extra_args:
         yue2_phase: "semantic"
         yue2_seed: 831001
         yue2_max_audio_frames: 200
+        yue2_max_hold_steps: 4096
       detokenize: false
-      # Engine-side backstop; the model stops on its own end token first.
-      max_tokens: 9000
+      stop_token_ids: [151848, 151852]  # ABC_END, MUSIC_END
+      # 200 frames + 4096 synthesis holds + the final end token. Callers
+      # changing the frame budget must also reserve this token headroom.
+      max_tokens: 4297
       temperature: 1.0

vllm_omni/entrypoints/openai/tts_adapters/yue2.py

@@ -36,6 +36,7 @@ from vllm_omni.model_executor.models.yue2.yue2 import (
     CONTEXT,
     FRAMES_PER_SECOND,
     KEY_MAX_AUDIO_FRAMES,
+    KEY_MAX_HOLD_STEPS,
     KEY_MIN_TOKENS,
     KEY_PENALTY_WINDOW,
     KEY_PHASE,
@@ -48,6 +49,7 @@ from vllm_omni.model_executor.models.yue2.yue2 import (
     KEY_TOP_P,
     SEMANTIC_SAMPLING,
     STOP_TOKEN_IDS,
+    SYNTHESIS_HOLD_STEPS,
 )
 from vllm_omni.tokenizers.yue2_prompt import semantic_prefix_ids
 
@@ -161,7 +163,9 @@ class Yue2Adapter(ARTTSAdapter):
         lyrics = request.input or ""
         prefix = semantic_prefix_ids(encode, instructions, lyrics, cot, abc_ids=abc_ids)
         max_frames = int(request.max_new_tokens or _DEFAULT_MAX_FRAMES)
-        if len(prefix) + max_frames > CONTEXT:
+        # +1: the song is delivered on the step after its last frame (the
+        # model acts on its draws one step late), which also needs a slot.
+        if len(prefix) + max_frames + 1 > CONTEXT:
             return_err = (
                 f"YuE2 prompt is {len(prefix)} tokens and the {max_frames}-frame budget "
                 f"exceeds the {CONTEXT}-token context. Shorten the lyrics/style or "
@@ -209,7 +213,12 @@ class Yue2Adapter(ARTTSAdapter):
         }
         # The model needs one decode step beyond the last frame to deliver the
         # terminal NAR/VAE pass, so do not cap the engine at exactly max_frames.
-        params.max_tokens = max_frames + 1
+        # Synthesize on a side stream while other songs keep decoding; the
+        # request may emit up to SYNTHESIS_HOLD_STEPS hold tokens meanwhile
+        # (see the model), so max_tokens leaves that much headroom.
+        prompt_len = len(prompt["prompt_token_ids"])
+        params.max_tokens = min(max_frames + 1 + SYNTHESIS_HOLD_STEPS, CONTEXT - prompt_len)
+        params.extra_args[KEY_MAX_HOLD_STEPS] = max(0, params.max_tokens - max_frames - 1)
         params.stop_token_ids = list(STOP_TOKEN_IDS)
         params.detokenize = False
         return sampling_params_list

vllm_omni/tokenizers/yue2_prompt.py

@@ -16,6 +16,7 @@ from vllm_omni.model_executor.models.yue2.yue2 import (
     ABC_START,
     CODEC_OFFSET,
     EOD,
+    HOLD_TOKEN,
     MUSIC_END,
     MUSIC_START,
 )
@@ -84,6 +85,10 @@ def semantic_frames(generated: list[int]) -> list[int]:
     ids = list(generated)
     if ids and ids[-1] == MUSIC_END:
         ids.pop()
+    # The model emits HOLD_TOKEN on the step it draws MUSIC_END (it acts on
+    # its draws one step late) and while the song is synthesized.
+    while ids and ids[-1] == HOLD_TOKEN:
+        ids.pop()
     return [t - CODEC_OFFSET for t in ids]
 
 

vllm_omni/worker/gpu_ar_model_runner.py

@@ -893,6 +893,44 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
             merged[req_id] = data
         return merged
 
+    def _call_prepare_runner_inputs(
+        self,
+        prepare_runner_inputs: Callable[..., tuple[torch.Tensor, torch.Tensor]],
+        *,
+        req_ids: list[str],
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None,
+        num_computed_tokens: np.ndarray,
+        num_scheduled_tokens: np.ndarray,
+        input_ids_buffer: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Keep opt-in request metadata outside captured model forwards."""
+        if getattr(self.model, "accepts_runner_sampling_extra_args", False):
+            return prepare_runner_inputs(
+                input_ids=input_ids,
+                positions=positions,
+                inputs_embeds=inputs_embeds,
+                req_ids=req_ids,
+                num_computed_tokens=num_computed_tokens,
+                num_scheduled_tokens=num_scheduled_tokens,
+                input_ids_buffer=input_ids_buffer,
+                sampling_extra_args=[
+                    params.extra_args if params and params.extra_args else {}
+                    for params in (self.requests[rid].sampling_params for rid in req_ids)
+                ],
+                discard_mask=self.discard_request_mask.np[: len(req_ids)],
+            )
+        return prepare_runner_inputs(
+            input_ids=input_ids,
+            positions=positions,
+            inputs_embeds=inputs_embeds,
+            req_ids=req_ids,
+            num_computed_tokens=num_computed_tokens,
+            num_scheduled_tokens=num_scheduled_tokens,
+            input_ids_buffer=input_ids_buffer,
+        )
+
     @torch.inference_mode()
     def execute_model(
         self,
@@ -1150,11 +1188,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         # for multimodal position detection, fix decode position offsets).
         prepare_runner_inputs = getattr(self.model, "prepare_runner_inputs", None)
         if callable(prepare_runner_inputs):
-            input_ids, positions = prepare_runner_inputs(
+            input_ids, positions = self._call_prepare_runner_inputs(
+                prepare_runner_inputs,
+                req_ids=req_ids[:num_reqs],
                 input_ids=input_ids,
                 positions=positions,
                 inputs_embeds=inputs_embeds,
-                req_ids=req_ids[:num_reqs],
                 num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs],
                 num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
                 input_ids_buffer=self.input_ids.gpu[:num_tokens_padded],