diff --git a/examples/offline_inference/yue2/end2end.py b/examples/offline_inference/yue2/end2end.py index 39fddd371..4d04a5633 100644 --- a/examples/offline_inference/yue2/end2end.py +++ b/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, diff --git a/recipes/m-a-p/YuE2-3B.md b/recipes/m-a-p/YuE2-3B.md index 8ad3b008d..197156e66 100644 --- a/recipes/m-a-p/YuE2-3B.md +++ b/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 diff --git a/tests/e2e/features/yue2/worker_extension.py b/tests/e2e/features/yue2/worker_extension.py new file mode 100644 index 000000000..ed3557560 --- /dev/null +++ b/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 diff --git a/tests/e2e/online_serving/test_yue2.py b/tests/e2e/online_serving/test_yue2.py index 35b25b1b0..17b020517 100644 --- a/tests/e2e/online_serving/test_yue2.py +++ b/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() diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 79d8f610d..daa8eabce 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/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"} diff --git a/tests/worker/test_yue2_request_state.py b/tests/worker/test_yue2_request_state.py new file mode 100644 index 000000000..75b68c8f3 --- /dev/null +++ b/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) diff --git a/tests/worker/test_yue2_sampler.py b/tests/worker/test_yue2_sampler.py new file mode 100644 index 000000000..4a72b0a13 --- /dev/null +++ b/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]) diff --git a/vllm_omni/deploy/yue2.yaml b/vllm_omni/deploy/yue2.yaml index 502d950c5..109212824 100644 --- a/vllm_omni/deploy/yue2.yaml +++ b/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 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/yue2.py b/vllm_omni/entrypoints/openai/tts_adapters/yue2.py index bcacc34e7..1d343afb8 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/yue2.py +++ b/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 diff --git a/vllm_omni/model_executor/models/yue2/yue2.py b/vllm_omni/model_executor/models/yue2/yue2.py index 02549cf84..bd1dd8b94 100644 --- a/vllm_omni/model_executor/models/yue2/yue2.py +++ b/vllm_omni/model_executor/models/yue2/yue2.py @@ -16,7 +16,7 @@ import json import math import os import time -from collections.abc import Callable, Iterable, Sequence +from collections.abc import Callable, Generator, Iterable, Iterator, Sequence from dataclasses import dataclass, field from numbers import Integral from pathlib import Path @@ -45,7 +45,8 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput Values mirror the upstream ``yue2/protocol.py`` (bd90e4c) exactly; changing any of them desynchronizes generation from the reference implementation. The model is single-codebook: one codec token per latent frame, 25 frames per second, -so ``max_tokens`` in engine terms equals frames/25 seconds of music. +so the frame budget divided by 25 is the maximum duration in seconds. +The engine token budget also includes synthesis HOLD and terminal tokens. """ EOD = 151643 @@ -59,6 +60,10 @@ VOCAB_SIZE, CONTEXT = 184704, 24576 # only ever emits the end token of the request's current phase, so a request # never stops on the other phase's end by accident. STOP_TOKEN_IDS = [ABC_END, MUSIC_END] +# Emitted in place of a freshly drawn MUSIC_END: the sampler reads its draws +# back one step late, so the request stays alive for one more step, in which +# the song is synthesized and MUSIC_END is emitted. +HOLD_TOKEN = LATENT_PAD SAMPLE_RATE = 48000 AUDIO_CHANNELS = 2 @@ -107,6 +112,15 @@ KEY_SKIP_SYNTHESIS = "yue2_skip_synthesis" # abc phase: tokens only, no audio # whole prefix, but under a KV prefix-cache hit the engine schedules only the # uncached tail, so the scheduled input_ids slice cannot rebuild it. KEY_PREFIX_IDS = "yue2_prefix_ids" +# How many steps a finishing song may emit HOLD_TOKEN while its NAR/VAE pass +# runs on a side stream (0: synthesize synchronously in the finishing step). +# The caller must leave that many tokens of max_tokens headroom. +KEY_MAX_HOLD_STEPS = "yue2_max_hold_steps" +# What the serving adapter and the offline example grant. A song holds while +# it waits for, and runs on, the side stream; under load the wait behind +# earlier songs dominates (about a thousand steps with 32-48 songs in flight +# on one H200). A request out of budget waits for its song in that step. +SYNTHESIS_HOLD_STEPS = 4096 # -------------------- checkpoint weight routing -------------------- @@ -156,80 +170,187 @@ request-local arithmetic exactly: phase masking (abc = text vocabulary + its end; semantic = codec span + its end), windowed repetition penalty over the request's own history, then temperature/top-k/top-p with a seeded multinomial. Default CFG is off (guidance 1.0 for full/melody, 1.01 for off); the upstream -torch backend treats those the same, so Phase 1 samples a single row. +torch backend treats those the same, so a request samples a single row. + +Rows sharing a preset are sampled as one batch over their phase's columns +only, as a CUDA graph per row count; the multinomial is drawn the way +``torch.multinomial`` draws one sample, so each request's random stream and +every drawn token match the per-row reference arithmetic. """ -def window_penalty(logits: torch.Tensor, recent_ids: list[int], penalty: float) -> torch.Tensor: - """Upweight/downweight ids seen in the window, upstream arithmetic.""" - if penalty == 1.0 or not recent_ids: - return logits - recent = torch.as_tensor(recent_ids, dtype=torch.long, device=logits.device) - if logits.dim() > 1: - recent = recent.reshape(1, -1) - freq = torch.zeros_like(logits) - freq.scatter_add_(-1, recent, torch.ones_like(recent, dtype=logits.dtype)) - alpha = penalty**freq - return torch.where(logits < 0, logits * alpha, logits / alpha) +# The semantic phase may only emit MUSIC_END or a codec token, and MUSIC_END +# sits right below the codec span, so its whole vocabulary is one slice. +assert MUSIC_END + 1 == CODEC_OFFSET -def distribution( +class PhaseVocab: + """Columns one phase may sample, in increasing token-id order. + + Phase masking sets every other column to -inf, so dropping them changes + no finite score, and keeping the id order keeps sort/argmax tie breaking + the same as on the full vocabulary. + """ + + def __init__(self, phase: str): + self.phase = phase + if phase == "abc": + self.end_col = EOD # columns [0, EOD) are ids; column EOD is ABC_END + self.num_cols = EOD + 1 + else: + self.end_col = 0 # column c is id MUSIC_END + c + self.num_cols = CODEC_OFFSET + CODEC_SIZE - MUSIC_END + self._index: dict[torch.device, torch.Tensor] = {} + + def take(self, rows: torch.Tensor) -> torch.Tensor: + """Allowed columns of ``rows`` [R, vocab] as a new float32 tensor.""" + if self.phase == "abc": + index = self._index.get(rows.device) + if index is None: + index = torch.cat((torch.arange(EOD), torch.tensor([ABC_END]))).to(rows.device) + self._index[rows.device] = index + return rows.index_select(-1, index).float() + return rows[..., MUSIC_END : CODEC_OFFSET + CODEC_SIZE].to(torch.float32, copy=True) + + def col(self, token_id: int) -> int: + if self.phase == "abc": + return self.end_col if token_id == ABC_END else token_id + return token_id - MUSIC_END + + def ids(self, cols: torch.Tensor) -> torch.Tensor: + if self.phase == "abc": + return torch.where(cols == self.end_col, ABC_END, cols) + return cols + MUSIC_END + + +PHASE_VOCAB = {"abc": PhaseVocab("abc"), "semantic": PhaseVocab("semantic")} + + +def _sample_core( logits: torch.Tensor, - *, + noise: torch.Tensor, + vocab: PhaseVocab, temperature: float, top_p: float, top_k: int, repetition_penalty: float, - penalty_window: int, - history: list[int], - step: int, - min_tokens: int, - phase: str, -) -> torch.Tensor: - """Masked/penalized/shaped scores for one row (float32 path). - - Accepts either one row ``[vocab]`` (what the model's sampler sees) or a - batch ``[rows, vocab]`` and returns the same shape it was given. - """ - single = logits.dim() == 1 - logits = logits.unsqueeze(0) if single else logits - scores = logits.float().clone() - end = ABC_END if phase == "abc" else MUSIC_END - allowed = torch.full_like(scores, float("-inf")) - if phase == "abc": - allowed[..., :EOD] = 0 - else: - allowed[..., CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0 - allowed[..., end] = 0 - scores = scores + allowed - if step < min_tokens: - scores[..., end] = -torch.inf - scores = window_penalty(scores, history[-penalty_window:], repetition_penalty) + window_cols: torch.Tensor | None, + block_end: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Masked, penalized, shaped scores and the multinomial draw, given the + Exp(1) noise [R, vocab].""" + rows = logits.shape[0] + scores = vocab.take(logits) + if block_end is not None: + scores[:, vocab.end_col].masked_fill_(block_end, -torch.inf) + if repetition_penalty != 1.0 and window_cols is not None: + # The padding column is counted, then dropped. + freq = torch.zeros((rows, vocab.num_cols + 1), dtype=scores.dtype, device=scores.device) + freq.scatter_add_(-1, window_cols, torch.ones_like(window_cols, dtype=scores.dtype)) + alpha = repetition_penalty ** freq[:, : vocab.num_cols] + scores = torch.where(scores < 0, scores * alpha, scores / alpha) + bad = ~torch.isfinite(scores).any(-1) if temperature == 0: - return scores.squeeze(0) if single else scores + return vocab.ids(scores.argmax(-1)), bad if temperature != 1: scores = scores / temperature - threshold = scores.topk(min(top_k, scores.shape[-1])).values[..., -1, None] + threshold = scores.topk(min(top_k, vocab.num_cols)).values[..., -1, None] scores = scores.masked_fill(scores < threshold, -torch.inf) if top_p < 1: - values, indices = scores.sort(descending=True) + values, indices = scores.sort(descending=True, stable=True) probabilities = values.softmax(-1) removed = probabilities.cumsum(-1) - probabilities > top_p removed[..., :1] = False values = values.masked_fill(removed, -torch.inf) scores = values.scatter(-1, indices, values) - return scores.squeeze(0) if single else scores - - -def sample_row(scores: torch.Tensor, generator: torch.Generator, *, greedy: bool = False) -> int: - """Draw one token id from prepared scores with the request's generator.""" - if not torch.isfinite(scores).any(): - raise RuntimeError("sampling scores are all -inf; check phase masking") - if greedy: - return int(scores.argmax().item()) probabilities = scores.softmax(-1) - next_id = torch.multinomial(probabilities, 1, generator=generator) - return int(next_id.item()) + return vocab.ids((probabilities / vocab.take(noise)).argmax(-1)), bad + + +def draw_noise(noise: torch.Tensor, generators: Sequence[torch.Generator]) -> torch.Tensor: + """Row r = the Exp(1) draw ``torch.multinomial`` makes from generator r.""" + for row, generator in enumerate(generators): + noise[row].exponential_(1, generator=generator) + return noise + + +def sample_rows( + logits: torch.Tensor, + vocab: PhaseVocab, + *, + temperature: float, + top_p: float, + top_k: int, + repetition_penalty: float, + window_cols: torch.Tensor | None, + block_end: torch.Tensor | None, + generators: list[torch.Generator], +) -> tuple[torch.Tensor, torch.Tensor]: + """Sample one token per row for rows sharing one preset. + + Works on the phase's allowed columns only and never synchronizes: every + input is already on the device (``window_cols`` [R, W] are the recent + ids as vocab columns, padded with ``vocab.num_cols``; ``block_end`` [R] + masks the end token) and it returns device tensors (token ids [R], + all-masked flags [R]). The multinomial draw is written out the way + ``torch.multinomial`` draws one sample (``argmax(p / q)`` with + ``q ~ Exp(1)`` over the full vocabulary), so each request's generator + advances exactly as in the per-row path. + """ + noise = torch.empty(logits.shape, dtype=torch.float32, device=logits.device) + draw_noise(noise, generators) + return _sample_core(logits, noise, vocab, temperature, top_p, top_k, repetition_penalty, window_cols, block_end) + + +class SamplerGraph: + """``sample_rows`` for a row-count bucket and preset, replayed as one graph. + + The sampler is ~30 small kernels per step, each a host-side launch; in a + batch-1 decode step that host time exceeds the backbone's GPU time. The + Exp(1) noise is still drawn per row from each request's own generator + outside the graph, and an all-False ``block_end`` / padded ``window_cols`` + leave every score as the eager path computes it, so draws are unchanged. + """ + + def __init__(self, vocab, rows, vocab_size, logits_dtype, preset, window, device, pool): + self.vocab = vocab + self.preset = preset # (temperature, top_p, top_k, repetition_penalty) + self.logits = torch.zeros((rows, vocab_size), dtype=logits_dtype, device=device) + self.noise = torch.ones((rows, vocab_size), dtype=torch.float32, device=device) + self.block_end = torch.zeros(rows, dtype=torch.bool, device=device) + self.window_cols = ( + torch.full((rows, window), vocab.num_cols, dtype=torch.long, device=device) if window else None + ) + side = torch.cuda.Stream(device) + side.wait_stream(torch.cuda.current_stream(device)) + with torch.cuda.stream(side): + for _ in range(2): + self._run() + torch.cuda.current_stream(device).wait_stream(side) + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph, pool=pool, capture_error_mode="thread_local"): + self.ids, self.bad = self._run() + + def _run(self): + return _sample_core(self.logits, self.noise, self.vocab, *self.preset, self.window_cols, self.block_end) + + def __call__(self, logits, window_cols, block_end, generators): + rows = logits.shape[0] + if rows > self.logits.shape[0] or len(generators) != rows: + raise ValueError("Sampler inputs exceed the captured bucket or lack a request generator") + # A replay reads its inputs from the addresses it was captured with. + # Padding rows have independent arithmetic and never advance a real + # request's generator. Only the active output rows are returned. + self.logits[:rows].copy_(logits) + draw_noise(self.noise, generators) + if self.window_cols is not None: + self.window_cols[:rows].copy_(window_cols) + if block_end is None: + self.block_end.zero_() + else: + self.block_end[:rows].copy_(block_end) + self.graph.replay() + return self.ids[:rows], self.bad[:rows] # -------------------- NAR flow-matching pass -------------------- @@ -238,11 +359,13 @@ def sample_row(scores: torch.Tensor, generator: torch.Generator, *, greedy: bool """Acoustic flow matching for YuE2-3B, ported from upstream ``nar.py``. One AR prefill per original chunk, then 32 midpoint ODE steps that walk only -the NAR path. The AR prefill re-uses the vLLM backbone's own projections -(fused ``qkv_proj`` split by head counts, ``q_norm``/``k_norm``, the layer's -``o_proj``/``mlp``), with RoPE and SDPA computed here exactly like the -reference: NAR positions attend the full AR prefix plus all NAR positions -bidirectionally, which a per-chunk prefill + concatenated K/V reproduces. +the NAR path: NAR positions attend the full AR prefix plus all NAR positions +bidirectionally. The prefill re-uses the vLLM backbone's own projections and +writes the prefix K/V once into a per-layer buffer; every ODE evaluation +writes its NAR K/V right after it and attends one contiguous sequence. An +evaluation is FlashAttention-3 (SDPA without it) plus the per-layer halves in +``_nar_pre_attention``/``_nar_post_attention`` (compiled), and after the +first call a CUDA graph replay: all 64 evaluations of a chunk share shapes. """ @@ -282,47 +405,158 @@ def _linear(module, x): return out[0] if isinstance(out, tuple) else out -def _attention(q, k, v, *, causal: bool, query_chunk_size: int | None = None): - """SDPA over [tokens, heads, dim]; query tiling only bounds temporaries.""" - if q.ndim != 3 or k.ndim != 3 or v.shape != k.shape or q.shape[-1] != k.shape[-1]: - raise ValueError("Expected Q/K/V [tokens, heads, dim] with matching K/V") - if min(q.shape) < 1 or min(k.shape) < 1 or q.shape[1] % k.shape[1]: - raise ValueError("Invalid attention lengths or grouped-query head count") - if causal and len(q) != len(k): - raise ValueError("Causal prefill requires matching Q/K sequence lengths") - block = query_chunk_size or (len(q) if q.device.type == "cuda" else 256) - query = q.transpose(0, 1).unsqueeze(0) - key = k.transpose(0, 1).unsqueeze(0) - value = v.transpose(0, 1).unsqueeze(0) - grouped = query.shape[1] != key.shape[1] - outputs = [] - for start in range(0, len(q), block): - end = min(start + block, len(q)) - used_key = key[..., :end, :] if causal else key - used_value = value[..., :end, :] if causal else value - mask = None - if causal and start: - mask = torch.arange(end, device=q.device)[None, :] <= torch.arange(start, end, device=q.device)[:, None] - outputs.append( - F.scaled_dot_product_attention( - query[..., start:end, :], - used_key, - used_value, - attn_mask=mask, - is_causal=causal and start == 0, - enable_gqa=grouped, +def _attention(q, k, v, *, causal: bool) -> torch.Tensor: + """SDPA over [tokens, heads, dim] Q/K/V (grouped K/V heads allowed).""" + out = F.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + k.transpose(0, 1).unsqueeze(0), + v.transpose(0, 1).unsqueeze(0), + is_causal=causal, + enable_gqa=q.shape[1] != k.shape[1], + ) + return out[0].transpose(0, 1) + + +def _drain(work: Iterator[None]) -> torch.Tensor: + """Run a work generator to completion and return its result.""" + while True: + try: + next(work) + except StopIteration as done: + return done.value + + +def _rms(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + # Same arithmetic as _RMSNorm.forward. + return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps).to(x.dtype) * weight + + +def _rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """_apply_rotary for [T, heads, dim] with [T, dim/2] cos/sin.""" + return _apply_rotary(x, cos[:, None, :], sin[:, None, :]) + + +def _nar_pre_attention(x, w_norm, w_qkv, w_q_norm, w_k_norm, cos, sin, eps: float, heads: int, kv_heads: int): + """input norm -> fused q/k/v projection -> q/k norm -> RoPE, for [T, hidden].""" + tokens = x.shape[0] + head_dim = w_q_norm.shape[0] + qkv = F.linear(_rms(x, w_norm, eps), w_qkv) + q, k, v = qkv.split([heads * head_dim, kv_heads * head_dim, kv_heads * head_dim], dim=-1) + q = _rope(_rms(q.reshape(tokens, heads, head_dim), w_q_norm, eps), cos, sin) + k = _rope(_rms(k.reshape(tokens, kv_heads, head_dim), w_k_norm, eps), cos, sin) + return q, k, v.reshape(tokens, kv_heads, head_dim) + + +def _nar_post_attention(x, attn, w_o, w_norm, w_gate_up, w_down, eps: float): + """o_proj + residual -> pre-MLP norm -> fused gate/up SiLU MLP + residual.""" + x = x + F.linear(attn.flatten(1), w_o) + gate, up = F.linear(_rms(x, w_norm, eps), w_gate_up).chunk(2, dim=-1) + return x + F.linear(F.silu(gate) * up, w_down) + + +_COMPILED: dict[str, Callable] = {} + + +def _nar_kernels() -> tuple[Callable, Callable]: + """The per-layer NAR halves, compiled once for any token count. + + Compilation fuses the RMSNorm/RoPE/SiLU elementwise chains that dominate + the eager NAR pass after its GEMMs and attention. + """ + if "pre" not in _COMPILED: + pre, post = _nar_pre_attention, _nar_post_attention + if torch.cuda.is_available(): + pre = torch.compile(pre, dynamic=True, fullgraph=True, options=_COMPILE_OPTIONS) + post = torch.compile(post, dynamic=True, fullgraph=True, options=_COMPILE_OPTIONS) + _COMPILED.update(pre=pre, post=post) + return _COMPILED["pre"], _COMPILED["post"] + + +def _flash_attention_fn() -> Callable | None: + """FlashAttention-3 varlen from vLLM on Hopper, else None (use SDPA).""" + if "fa3" not in _COMPILED: + fn = None + if torch.cuda.is_available(): + try: + from vllm.vllm_flash_attn import flash_attn_interface, flash_attn_varlen_func + except ImportError: + pass + else: + if flash_attn_interface.is_fa_version_supported(3): + fn = flash_attn_varlen_func + _COMPILED["fa3"] = fn + return _COMPILED["fa3"] + + +def _sequence_bounds(length: int, device: torch.device) -> torch.Tensor: + return torch.tensor([0, length], dtype=torch.int32).to(device, non_blocking=True) + + +def _flash_attention(fa3, q, k, v, *, causal: bool, bounds_q=None, bounds_k=None) -> torch.Tensor: + """One sequence of [tokens, heads, dim] Q/K/V (grouped K/V heads allowed). + + ``fa3`` is FlashAttention-3 varlen, or None for PyTorch SDPA. + """ + if fa3 is None: + return _attention(q, k, v, causal=causal) + bounds_q = _sequence_bounds(q.shape[0], q.device) if bounds_q is None else bounds_q + bounds_k = _sequence_bounds(k.shape[0], q.device) if bounds_k is None else bounds_k + return fa3( + q, + k, + v, + max_seqlen_q=q.shape[0], + cu_seqlens_q=bounds_q, + max_seqlen_k=k.shape[0], + cu_seqlens_k=bounds_k, + causal=causal, + fa_version=3, + ) + + +def fuse_nar_projections(layers: nn.ModuleList) -> list[tuple[torch.Tensor, ...]]: + """Per NAR layer: norm/qkv/q-norm/k-norm/o/norm/gate-up/down weights. + + q/k/v and gate/up are concatenated so each runs as one GEMM; the original + ``nn.Linear`` weights become views into the fused tensors, so the module + tree stays intact at no extra memory. + """ + fused = [] + for layer in layers: + attn, mlp = layer.self_attn, layer.mlp + projections = (attn.q_proj, attn.k_proj, attn.v_proj) + qkv = torch.cat([p.weight.data for p in projections]) + offset = 0 + for p in projections: + rows = p.weight.shape[0] + p.weight.data = qkv[offset : offset + rows] + offset += rows + gate_up = torch.cat([mlp.gate_proj.weight.data, mlp.up_proj.weight.data]) + rows = mlp.gate_proj.weight.shape[0] + mlp.gate_proj.weight.data = gate_up[:rows] + mlp.up_proj.weight.data = gate_up[rows:] + fused.append( + ( + layer.input_layernorm.weight, + qkv, + attn.q_norm.weight, + attn.k_norm.weight, + attn.o_proj.weight, + layer.pre_mlp_layernorm.weight, + gate_up, + mlp.down_proj.weight, ) ) - return torch.cat(outputs, dim=-2)[0].transpose(0, 1) + return fused class CachedNAR: """One acoustic chunk: AR prefix K/V prefilled once, ODE walks the NAR path.""" - def __init__(self, model, chunk: Chunk, query_chunk_size: int | None = None): + def __init__(self, model, chunk: Chunk, pool: tuple[int, int] | None = None): self.model = model self.chunk = chunk - self.query_chunk_size = query_chunk_size + self.pool = pool weight = next(model.vae2llm.parameters()) self.device, self.dtype = weight.device, weight.dtype if chunk.noise.ndim != 2 or chunk.noise.shape[1] != 64 or len(chunk.noise) < 1: @@ -333,65 +567,182 @@ class CachedNAR: if self.ar_length + self.nar_length > model.max_position_embeddings: raise ValueError("Acoustic chunk exceeds the model context") positions = torch.arange(self.ar_length, self.ar_length + self.nar_length, device=self.device)[None] - self.cos, self.sin = model.rotary(positions) + cos, sin = model.rotary(positions) + self.cos, self.sin = cos[0], sin[0] # [NAR, dim/2] local = torch.arange(self.nar_length, device=self.device).clamp(max=model.max_latent_frames - 1) self.pos_emb = model.latent_pos_embed(local)[None] - self.cache: list[tuple[torch.Tensor, torch.Tensor]] = [] + self.fa3 = _flash_attention_fn() + self.bounds_q = _sequence_bounds(self.nar_length, self.device) + self.bounds_k = _sequence_bounds(self.ar_length + self.nar_length, self.device) + self.graph: torch.cuda.CUDAGraph | None = None + self.last_use: torch.cuda.Event | None = None self._prefill() def _prefill(self): - backbone = self.model.model - ids = torch.tensor([self.chunk.ar_tokens], dtype=torch.long, device=self.device) - positions = torch.arange(self.ar_length, device=self.device)[None] - cos, sin = self.model.rotary(positions) - x = backbone.embed_tokens(ids)[0] # [T, H] - for layer in backbone.layers: - # Upstream: layer.self_attn.project_qkv(layer.input_layernorm(x), ...) + model = self.model + backbone = model.model + config = model.config + total = self.ar_length + self.nar_length + shape = (len(backbone.layers), total, int(config.num_key_value_heads), int(config.head_dim)) + self.keys = torch.empty(shape, dtype=self.dtype, device=self.device) + self.values = torch.empty(shape, dtype=self.dtype, device=self.device) + ids = torch.tensor(self.chunk.ar_tokens, dtype=torch.long).to(self.device, non_blocking=True) + cos, sin = model.rotary(torch.arange(self.ar_length, device=self.device)[None]) + x = backbone.embed_tokens(ids) # [T, H] + last = len(backbone.layers) - 1 + for index, layer in enumerate(backbone.layers): # The layernorm is load-bearing: without it the cached AR K/V # conditioning is garbage and the ODE latents blow up (clipped audio). - q, k, v = self.model.ar_project(layer, layer.input_layernorm(x), cos, sin) - self.cache.append((k, v)) - h = _attention(q, k, v, causal=True, query_chunk_size=self.query_chunk_size) + q, k, v = model.ar_project(layer, layer.input_layernorm(x), cos, sin) + self.keys[index, : self.ar_length] = k + self.values[index, : self.ar_length] = v + if index == last: + break # the last layer's output feeds nothing + h = _flash_attention(self.fa3, q, k, v, causal=True) x = x + _linear(layer.self_attn.o_proj, h.flatten(1)) x = x + _linear(layer.mlp, layer.post_attention_layernorm(x)) @torch.inference_mode() def velocity(self, state: torch.Tensor, raw_t: float) -> torch.Tensor: + """One ODE evaluation; from the second call on, a CUDA graph replay. + + After one eager call (compilation, cuBLAS/FA3 setup) the evaluation is + captured once and replayed from static state/time inputs, replacing + ~30 host-side launches per layer with one launch per evaluation. + """ + if self.device.type != "cuda": + return self._velocity(state, torch.tensor(raw_t, dtype=self.dtype, device=self.device)) + if self.graph is None: + if not hasattr(self, "state_in"): + self.state_in = torch.empty_like(state) + self.time_in = torch.empty((), dtype=self.dtype, device=self.device) + self.state_in.copy_(state) + self.time_in.fill_(raw_t) + return self._velocity(self.state_in, self.time_in) + self.graph = torch.cuda.CUDAGraph() + if torch.cuda.current_stream(self.device) == torch.cuda.default_stream(self.device): + with torch.cuda.graph(self.graph, pool=self.pool, capture_error_mode="thread_local"): + self.velocity_out = self._velocity(self.state_in, self.time_in) + else: + # Already on a side stream (the async synthesis path): capture + # in place, without torch.cuda.graph's device-wide sync, which + # would stall the decode loop behind every queued kernel. + self.graph.capture_begin(pool=self.pool, capture_error_mode="thread_local") + try: + self.velocity_out = self._velocity(self.state_in, self.time_in) + finally: + self.graph.capture_end() + # The graph reads its inputs from the addresses it was captured with. + self.state_in.copy_(state) + self.time_in.fill_(raw_t) + self.graph.replay() + return self.velocity_out + + def _velocity(self, state: torch.Tensor, raw_t: torch.Tensor) -> torch.Tensor: model = self.model + pre, post = _nar_kernels() + config = model.config + eps = float(config.rms_norm_eps) + heads, kv_heads = int(config.num_attention_heads), int(config.num_key_value_heads) x_nar = F.pad(state, (0, 0, 1, 1)) - shifted = model.shift_t(raw_t, self.device, self.dtype) + # The timestep shift, from a device scalar so a graph replay can update it. + t_sig = torch.sigmoid(raw_t) + shift = model._timestep_shift + shifted = shift * t_sig / (1 + (shift - 1) * t_sig) x = model.vae2llm(x_nar[None]) x = x + model.time_embedder(shifted.expand(self.nar_length))[None] - x = x + self.pos_emb - for layer, (ar_k, ar_v) in zip(model.nar_layers, self.cache): - q, k, v = layer.self_attn.project_qkv(layer.input_layernorm(x), self.cos, self.sin) - k, v = torch.cat((ar_k, k[0])), torch.cat((ar_v, v[0])) - h = _attention(q[0], k, v, causal=False, query_chunk_size=self.query_chunk_size) - x = x + _linear(layer.self_attn.o_proj, h.flatten(1)[None]) - x = x + layer.mlp(layer.pre_mlp_layernorm(x)) - return model.llm2vae(model.model.norm(x))[0, 1:-1] + x = (x + self.pos_emb)[0] + for index, weights in enumerate(model.nar_weights()): + w_norm, w_qkv, w_q_norm, w_k_norm, w_o, w_mlp_norm, w_gate_up, w_down = weights + q, k, v = pre(x, w_norm, w_qkv, w_q_norm, w_k_norm, self.cos, self.sin, eps, heads, kv_heads) + self.keys[index, self.ar_length :] = k + self.values[index, self.ar_length :] = v + h = _flash_attention( + self.fa3, + q, + self.keys[index], + self.values[index], + causal=False, + bounds_q=self.bounds_q, + bounds_k=self.bounds_k, + ) + x = post(x, h, w_o, w_mlp_norm, w_gate_up, w_down, eps) + return model.llm2vae(model.model.norm(x[None]))[0, 1:-1] - @torch.inference_mode() - def solve(self, steps: int = 32) -> torch.Tensor: + def solve_iter(self, steps: int = 32) -> Iterator[None]: + """The ODE solve as a work generator: yields after enqueuing each + evaluation and returns the (still in flight) device latents.""" if isinstance(steps, bool) or not isinstance(steps, Integral) or steps < 1: raise ValueError("steps must be a positive integer") - state = self.chunk.noise.to(device=self.device, dtype=self.dtype) + # Non-blocking: a blocking copy would wait for the queued prefill. + state = self.chunk.noise.to(device=self.device, dtype=self.dtype, non_blocking=True) dt = 1.0 / steps for step in range(steps): t = 1.0 - step * dt raw = torch.logit(torch.tensor(t, dtype=torch.float64, device="cpu")).clamp(-20, 20).item() first = self.velocity(state, raw) + yield mid = state - first * (dt / 2) raw_mid = torch.logit(torch.tensor(t - dt / 2, dtype=torch.float64, device="cpu")).clamp(-20, 20).item() state = state - self.velocity(mid, raw_mid) * dt - result = state.float().cpu() + yield + return state.float() + + @torch.inference_mode() + def solve(self, steps: int = 32) -> torch.Tensor: + result = _drain(self.solve_iter(steps)).cpu() if not torch.isfinite(result).all(): raise FloatingPointError("Acoustic flow matching produced non-finite latents") return result def close(self): - self.cache.clear() + self.keys = self.values = self.graph = None self.cos = self.sin = self.pos_emb = None + self.state_in = self.time_in = self.velocity_out = None + self.bounds_q = self.bounds_k = None + self.chunk = None + + +class _NARGraphPool: + """Allocator pool shared by graphs on the serialized synthesis stream.""" + + def __init__(self) -> None: + self.pool: tuple[int, int] | None = None + self._anchor: torch.cuda.CUDAGraph | None = None + self._anchor_buffer: torch.Tensor | None = None + + def handle(self, device: torch.device) -> tuple[int, int]: + if self.pool is None: + self.pool = torch.cuda.graph_pool_handle() + # The handle alone does not keep the allocator pool alive. A tiny + # graph owns it after a chunk is released, including across aborts. + self._anchor = torch.cuda.CUDAGraph() + self._anchor.capture_begin(pool=self.pool, capture_error_mode="thread_local") + try: + self._anchor_buffer = torch.zeros((), dtype=torch.int32, device=device) + finally: + self._anchor.capture_end() + return self.pool + + +def synthesize_iter( + model, prefix: Sequence[int], codec: Sequence[int], seed: int, keep: list, graph_pool: _NARGraphPool | None = None +) -> Iterator[None]: + """``synthesize`` as a work generator (one unit per prefill or ODE + evaluation); returns the in-flight device latents, unchecked. The solver + engines go to ``keep``: their graphs and buffers must outlive the work.""" + output = [] + for chunk in song_chunks(prefix, codec, seed): + # Graphs share allocator storage only on this serialized stream. + # Their state and K/V are retired after each chunk, with no engine cache. + pool = None if graph_pool is None else graph_pool.handle(next(model.vae2llm.parameters()).device) + engine = CachedNAR(model, chunk, pool) + keep.append(engine) + yield + output.append((yield from engine.solve_iter(ODE_STEPS))) + engine.last_use = torch.cuda.Event() + engine.last_use.record() + return torch.cat(output, dim=0) @torch.inference_mode() @@ -440,6 +791,26 @@ def snake_beta(x, alpha, beta): return x + (1.0 / (beta + 0.000000001)) * torch.pow(torch.sin(x * alpha), 2) +_SNAKE: list[Callable] = [] +# Inductor otherwise benchmarks a few launch configs per reduction kernel and +# keeps the fastest; timing noise then changes the reduction order, and with it +# the song's last bits, from one process to the next. Deterministic mode keeps +# "same seed, same song" across server restarts. +_COMPILE_OPTIONS = {"deterministic": True} + + +def _snake_kernel() -> Callable: + """snake_beta fused into one kernel: it runs on every [C, samples] FP32 + activation of the decoder, where each unfused elementwise op is a full + memory pass.""" + if not _SNAKE: + fused = torch.cuda.is_available() + _SNAKE.append( + torch.compile(snake_beta, dynamic=True, fullgraph=True, options=_COMPILE_OPTIONS) if fused else snake_beta + ) + return _SNAKE[0] + + class SnakeBeta(nn.Module): def __init__( self, @@ -467,6 +838,8 @@ class SnakeBeta(nn.Module): if self.alpha_logscale: alpha = torch.exp(alpha) beta = torch.exp(beta) + if x.is_cuda and not self.training: + return _snake_kernel()(x, alpha, beta) return snake_beta(x, alpha, beta) @@ -981,11 +1354,13 @@ class YuE2VAE(PreTrainedModel): def decoder_device(self): return next(self.decoder.parameters()).device - def _latent(self, latent): + def _latent(self, latent, check_finite=True): latent = torch.as_tensor(latent) if latent.ndim != 3 or latent.shape[1] != self.config.latent_dim or latent.shape[0] < 1 or latent.shape[-1] < 1: raise ValueError(f"Expected nonempty [B,{self.config.latent_dim},T] latents") - if not torch.isfinite(latent).all(): + # The check reads the values back, a device sync; an async caller + # checks the decoded audio on the host instead. + if check_finite and not torch.isfinite(latent).all(): raise ValueError("VAE latents contain non-finite values") if next(self.decoder.parameters()).dtype != torch.float32: raise ValueError("VAE decoder weights must remain FP32") @@ -1024,9 +1399,9 @@ class YuE2VAE(PreTrainedModel): return latent @torch.inference_mode() - def decode(self, latent): + def decode(self, latent, check_finite=True): """Full waveform, FP32 [B,2,1920*T-64], without clipping.""" - latent = self._latent(latent) + latent = self._latent(latent, check_finite) with torch.autocast(device_type=self.decoder_device.type, enabled=False): return self.decoder(latent.to(device=self.decoder_device, dtype=torch.float32)) @@ -1049,6 +1424,20 @@ class YuE2VAE(PreTrainedModel): halo_frames=None, output_device="cpu", on_progress: Callable[[int, int], None] | None = None, + check_finite=True, + ): + return _drain( + self.decode_tiled_iter(latent, core_frames, halo_frames, output_device, on_progress, check_finite) + ) + + def decode_tiled_iter( + self, + latent, + core_frames=None, + halo_frames=None, + output_device="cpu", + on_progress: Callable[[int, int], None] | None = None, + check_finite=True, ): """Decode bounded tiles, retaining exact cores with natural end length. @@ -1059,7 +1448,7 @@ class YuE2VAE(PreTrainedModel): no extra synchronization is added. With a CUDA output device, queued work may still be executing. Callback exceptions propagate. """ - latent = self._latent(latent) + latent = self._latent(latent, check_finite) core_frames = self.config.decode_core_frames if core_frames is None else core_frames halo_frames = self.config.decode_halo_frames if halo_frames is None else halo_frames if not isinstance(core_frames, int) or core_frames < 1: @@ -1078,7 +1467,7 @@ class YuE2VAE(PreTrainedModel): end = min(frames, start + core_frames) left = max(0, start - halo_frames) right = min(frames, end + halo_frames) - tile = self.decode(latent[..., left:right]) + tile = self.decode(latent[..., left:right], check_finite) out_start, out_end = start * ratio, min(end * ratio, total) crop_start = (start - left) * ratio crop = tile[..., crop_start : crop_start + out_end - out_start] @@ -1088,6 +1477,7 @@ class YuE2VAE(PreTrainedModel): del tile, crop if on_progress is not None: on_progress(tile_index + 1, tiles) + yield return audio def decode_audio(self, latent, chunked=True, **kwargs): @@ -1180,15 +1570,6 @@ class _NARAttention(nn.Module): self.q_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) self.k_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) - def project_qkv(self, x, cos, sin): - B, T, _ = x.shape - q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim) - k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim) - v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim) - q, k = self.q_norm(q), self.k_norm(k) - rc, rs = cos.unsqueeze(2), sin.unsqueeze(2) - return _apply_rotary(q, rc, rs), _apply_rotary(k, rc, rs), v - class _NARMLP(nn.Module): def __init__(self, config): @@ -1265,6 +1646,7 @@ class _RowConstants: max_audio_frames: int seed: int skip_synthesis: bool + max_hold_steps: int = 0 @dataclass @@ -1273,13 +1655,263 @@ class _RequestState: request_id: str constants: _RowConstants - prompt_tokens: int = 0 prompt_len: int = 0 # full prompt length; a completing prefill row reaches it prefix_ids: list[int] = field(default_factory=list) history: list[int] = field(default_factory=list) generator: torch.Generator | None = None finished: bool = False truncated: bool = False + # Set when the host learns (one step after the draw) that the request + # drew its end token or filled its frame budget; the next scheduled step + # runs the finishing pass and emits the end token. + finish_ready: bool = False + end_drawn: bool = False # the engine counts a drawn end (or its HOLD) too + job: _SynthesisJob | None = None # async synthesis in flight + hold_steps: int = 0 # HOLD_TOKEN steps emitted while it runs + # Engine output count once the emitted end token is accepted, and the + # shipped (audio, truncated, error). A preemption can drop the step that + # delivered both; the resumed request then delivers them again. + end_len: int | None = None + delivered: tuple[torch.Tensor, bool, bool] | None = None + + +SongWork = Callable[[list], Generator[None, None, torch.Tensor]] + + +class _SynthesisJob: + """One song's NAR + VAE pass on a side stream, fed a little per step. + + The pass is a generator of work units (the prefill, each ODE evaluation, + each VAE tile). ``advance`` enqueues units only while fewer than + ``MAX_INFLIGHT`` are queued ahead of the GPU and never waits: enqueuing + the whole song at once fills the stream's launch queue, after which every + launch blocks the host until the GPU drains, and the decode loop with it. + The decode steps call ``advance`` every step, so the side stream stays + busy while the other requests keep decoding; kernels of the two streams + run concurrently within the process. The finishing request emits + HOLD_TOKEN until ``done()``. + """ + + MAX_INFLIGHT = 2 + + def __init__(self, request_id: str, song: SongWork): + """``song(keep)`` yields once per enqueued work unit and returns the + device audio [2, samples]; what it appends to ``keep`` stays alive + until the job is released.""" + self.request_id = request_id + self.song = song + self.stream: torch.cuda.Stream | None = None + self.work: Iterator[None] | None = None + self.inflight: list[torch.cuda.Event] = [] + self.event: torch.cuda.Event | None = None # recorded after the last unit + self.audio: torch.Tensor | None = None # pinned host copy + self.finite: torch.Tensor | None = None # pinned pre-clamp validity flag + self.error: BaseException | None = None + self.cancelled = False + self._released = False + # Solver engines (captured graphs, K/V buffers), latents and audio: + # they must outlive the queued kernels that use them. + self.keep: list[CachedNAR | torch.Tensor] = [] + self.t_start: float | None = None + self.t_done: float | None = None + + @property + def started(self) -> bool: + return self.t_start is not None + + def _run(self) -> Iterator[None]: + audio = yield from self.song(self.keep) + # Check before clipping: clipping maps infinities to finite +/-1. + # Both copies target pinned allocations and complete before self.event. + self.finite = torch.empty((), dtype=torch.bool, pin_memory=True) + self.finite.copy_(torch.isfinite(audio).all(), non_blocking=True) + audio = audio.clamp(-1, 1) + # Pinned, so the device-to-host copy stays async on the side stream. + self.audio = torch.empty(audio.shape, dtype=audio.dtype, pin_memory=True) + self.audio.copy_(audio, non_blocking=True) + self.keep.append(audio) + + def start(self, stream: torch.cuda.Stream) -> None: + self.stream = stream + self.t_start = time.perf_counter() + self.work = self._run() + self.advance() + + def advance(self) -> None: + """Enqueue work units up to MAX_INFLIGHT ahead of the GPU; never waits.""" + if self.work is None or self.event is not None or self.error is not None: + return + with torch.cuda.stream(self.stream), torch.inference_mode(): + while True: + # Completed chunks need not keep all their K/V and private + # graph pools alive until the whole song finishes. + for item in self.keep[:]: + if isinstance(item, CachedNAR) and item.last_use is not None and item.last_use.query(): + item.close() + self.keep.remove(item) + while self.inflight and self.inflight[0].query(): + self.inflight.pop(0) + if len(self.inflight) >= self.MAX_INFLIGHT: + return + try: + next(self.work) + except StopIteration: + self.event = torch.cuda.Event() + self.event.record(self.stream) + self.inflight.clear() + return + except (RuntimeError, ValueError) as exc: + # CUDA errors and OOM are RuntimeErrors; song_chunks + # rejects a bad song with ValueError. Either fails only + # this request (see finish); anything else is a bug. + self.error = exc + return + unit = torch.cuda.Event() + unit.record(self.stream) + self.inflight.append(unit) + + def done(self) -> bool: + if self.cancelled or self.error is not None: + return True + self.advance() + finished = self.event is not None and self.event.query() + if finished and self.t_done is None: + self.t_done = time.perf_counter() + return finished + + def drain(self) -> None: + """Enqueue and wait for everything (a request out of hold budget).""" + while self.work is not None and self.event is None and self.error is None: + if self.inflight: + self.inflight[0].synchronize() + self.advance() + + def release(self) -> None: + """Release once, after this job's last submitted kernel has finished. + + The queue may already be running a later job on the same stream when + this job is delivered. Waiting on the stream would wait for that job. + """ + if self._released: + return + if self.stream is not None and self.event is None: + # An error or cancellation can precede the normal terminal event. + self.event = torch.cuda.Event() + self.event.record(self.stream) + if self.event is not None: + self.event.synchronize() + if self.work is not None: + self.work.close() + self.work = None + for item in self.keep: + if isinstance(item, CachedNAR): + item.close() + self.keep.clear() + self.stream = None + self._released = True + + def cancel(self) -> None: + """Stop submitting work and release the already submitted units. + + Waiting is bounded by MAX_INFLIGHT, not the rest of the song. Cleanup + happens here even if the engine becomes idle after the cancellation. + """ + self.cancelled = True + self.release() + + def finish(self) -> tuple[torch.Tensor, bool]: + """Wait if needed, release, and return (audio [2, T] on the CPU, failed).""" + self.drain() + self.release() + if self.t_done is None: + self.t_done = time.perf_counter() + if self.cancelled: + return torch.zeros((2, 0)), True + if self.error is not None: + logger.error("YuE2 synthesis failed for req=%s: %r", self.request_id, self.error) + return torch.zeros((2, 0)), True + if self.finite is None or not bool(self.finite): + logger.error("YuE2 synthesis produced non-finite audio for req=%s", self.request_id) + return torch.zeros((2, 0)), True + return self.audio, False + + +class _SynthesisQueue: + """Songs waiting for, or running on, the synthesis side stream. + + One song runs at a time: the NAR pass is compute-bound, so a second one + would only double the memory. ``pump`` runs every decode step; it feeds + the running song and starts the next once it is done. + """ + + def __init__(self) -> None: + self._stream: torch.cuda.Stream | None = None + self.graph_pool = _NARGraphPool() + self.active: _SynthesisJob | None = None + self.waiting: list[_SynthesisJob] = [] + + def stream(self) -> torch.cuda.Stream: + """The side stream, at the highest priority the device offers. + + Once several songs are in flight the synthesis pass, not the decode, + bounds throughput, and its kernels leave no room to co-schedule the + decode's anyway: they fill every SM. So synthesis goes first and the + decode fills the gaps. + """ + if self._stream is None: + self._stream = torch.cuda.Stream(priority=torch.cuda.Stream.priority_range()[1]) + return self._stream + + def submit(self, job: _SynthesisJob) -> None: + self.waiting.append(job) + self.pump() + + def pump(self) -> None: + """Feed the running song; start the next queued one once it is done.""" + if self.active is not None: + if not self.active.done(): # done() also feeds it + return + self.active.release() # frees its buffers; the audio stays + self.active = None + while self.waiting: + job = self.waiting.pop(0) + if job.error is not None: + continue # fails at completion, without touching the GPU + job.start(self.stream()) + self.active = job + return + + def complete(self, job: _SynthesisJob) -> tuple[torch.Tensor, bool, bool]: + """Finish ``job`` now, ahead of the queue if it has not started. + + Returns (audio [2, T] on the CPU, failed, whether it had to wait). + """ + if job.cancelled: + return torch.zeros((2, 0)), True, False + if job in self.waiting: + self.waiting.remove(job) + if not job.started and job.error is None: + if self.active is not None: + # Finish the running song first (its own request ships it). + self.active.drain() + self.active.release() + self.active = None + job.start(self.stream()) + self.active = job + waited = not job.done() + audio, failed = job.finish() + if self.active is job: + self.active = None + self.pump() + return audio, failed, waited + + def cancel(self, job: _SynthesisJob) -> None: + """Drop an aborted song and release its bounded in-flight work.""" + if job in self.waiting: + self.waiting.remove(job) + job.cancel() + if self.active is job: + self.active = None class Yue2ForCausalLM(nn.Module): @@ -1288,6 +1920,9 @@ class Yue2ForCausalLM(nn.Module): have_multimodal_outputs = True prefer_model_sampler = True has_postprocess = False + # Request state is created in prepare_runner_inputs, which runs on every + # step; forward() is skipped on a full CUDA graph replay. + accepts_runner_sampling_extra_args = True # The song is decoded in-model and shipped via make_omni_output; per-step # hidden states must NOT ride the pooler payload, or the output pipeline # remaps the accumulated "hidden" rows to the audio modality key and the @@ -1328,7 +1963,6 @@ class Yue2ForCausalLM(nn.Module): prefix=maybe_prefix(prefix, "lm_head"), ) self.logits_processor = LogitsProcessor(int(config.vocab_size)) - self.nar_layers = nn.ModuleList([_NARLayer(config) for _ in range(int(config.num_hidden_layers))]) self.vae2llm = nn.Linear(int(getattr(config, "latent_dim", LATENT_DIM)), hidden) self.llm2vae = nn.Linear(hidden, int(getattr(config, "latent_dim", LATENT_DIM))) @@ -1343,27 +1977,31 @@ class Yue2ForCausalLM(nn.Module): self.max_position_embeddings = int(config.max_position_embeddings) self.max_latent_frames = int(getattr(config, "max_latent_frames", self.max_position_embeddings)) self._timestep_shift = float(getattr(config, "timestep_shift", 1.0)) - - self._states: dict[str, _RequestState] = {} - self._row_constants: dict[str, _RowConstants] = {} - self._audio_queue: list[tuple[str, torch.Tensor, bool, bool]] = [] - self._deferred_cleanup_ids: set[str] = set() - self._step_rows: list[tuple[str, int, int]] = [] # (req_id, computed, scheduled) - self._decode_t0: dict[str, float] = {} # first decode-step wall time, for AR tok/s - self._last_mm: dict[str, Any] | None = None # current step's make_omni_output payload + # Fused NAR projection weights (fuse_nar_projections); see nar_weights. + self.nar_fused: list[tuple[torch.Tensor, ...]] | None = None self._vae: YuE2VAE | None = None self._vae_device: torch.device | None = None + self._states: dict[str, _RequestState] = {} + self._audio_queue: list[tuple[str, torch.Tensor, bool, bool]] = [] + self._step_rows: list[tuple[str, int, int]] = [] # (req_id, computed, scheduled) + self._step_discard: list[bool] | None = None + self._decode_t0: dict[str, float] = {} # first decode-step wall time, for AR tok/s + self._last_mm: dict[str, Any] | None = None # current step's make_omni_output payload + # Tokens drawn by the last sample() call, still on their way to the + # host: (copy-done event, pinned [2, rows] ids/flags, states). + self._pending: tuple[torch.Event, torch.Tensor, list[_RequestState]] | None = None + self._pinned: torch.Tensor | None = None + self._stage: torch.Tensor | None = None + self._logits: torch.Tensor | None = None # semantic-only logits buffer + self._sampler_graphs: dict[tuple, SamplerGraph] = {} + self._sampler_graph_pool: tuple[int, int] | None = None + self._sampler_capacity = int(vllm_config.scheduler_config.max_num_seqs) + self._max_model_len = int(vllm_config.model_config.max_model_len) + self._synthesis = _SynthesisQueue() + # ------------------------------------------------------------ sampling - def shift_t(self, raw_t: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor: - t_sig = torch.sigmoid(torch.tensor(raw_t, dtype=dtype, device=device)) - shift = self._timestep_shift - return shift * t_sig / (1 + (shift - 1) * t_sig) - - def rotary(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - return self.rotary_emb(positions) - def ar_project(self, layer, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): """AR-path Q/K/V from the vLLM backbone's fused projection.""" attn = layer.self_attn @@ -1387,6 +2025,108 @@ class Yue2ForCausalLM(nn.Module): k = _apply_rotary(k.unsqueeze(0), rc, rs).squeeze(0) return q, k, v.view(T, num_kv_heads, head_dim) + def nar_weights(self) -> list[tuple[torch.Tensor, ...]]: + """The fused NAR layer weights, fused on first use. + + load_weights fuses them, but a dummy-weight load (load_format=dummy, + as the CI smoke tests run) never calls it. + """ + if self.nar_fused is None: + self.nar_fused = fuse_nar_projections(self.nar_layers) + return self.nar_fused + + def rotary(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return self.rotary_emb(positions) + + def synthesize_song( + self, prefix_ids: Sequence[int], codec: Sequence[int], seed: int + ) -> tuple[torch.Tensor, float, float]: + """Channels-first stereo audio for one song, plus NAR and VAE seconds.""" + t0 = time.perf_counter() + latents = synthesize(self, prefix_ids, codec, seed, steps=ODE_STEPS) + logger.debug( + "YUE2_LATENTS frames=%d absmax=%.4f mean=%.5f std=%.4f", + latents.shape[0], + float(latents.abs().max()), + float(latents.mean()), + float(latents.std()), + ) + t1 = time.perf_counter() + audio = self._decode_latents(latents) + return audio, t1 - t0, time.perf_counter() - t1 + + @torch.inference_mode() + def _warm_up_synthesis(self) -> None: + """Compile the NAR kernels and touch the VAE once at startup. + + Otherwise the first finishing song pays the compile (seconds) inside + a decode step. One ODE step over a few frames covers every shape: + the NAR kernels are compiled for dynamic token counts. + """ + if not torch.cuda.is_available() or self.nar_fused is None: + return + t0 = time.perf_counter() + training, default_dtype = self.training, torch.get_default_dtype() + # load_weights runs before the loader switches to eval, and under the + # loader's bf16 default dtype, which the compiled kernels guard on: + # warm up in the serving-time state or the first song recompiles. + self.eval() + torch.set_default_dtype(torch.float32) + try: + latents = synthesize(self, [EOD, ABC_START, ABC_END, MUSIC_START], list(range(16)), DEFAULT_SEED, steps=1) + self._decode_latents(latents) + finally: + torch.set_default_dtype(default_dtype) + self.train(training) + logger.info("YuE2 synthesis warm-up took %.1f s", time.perf_counter() - t0) + + def _vae_model(self, device: torch.device) -> YuE2VAE: + if self._vae is None: + vae_path = os.environ.get("YUE2_VAE", DEFAULT_VAE_ID) + logger.info("Loading YuE2 VAE decoder from %s", vae_path) + # Keep the VAE out of the module tree (__dict__, not an nn.Module + # attribute): the checkpoint has no VAE keys, so a registered + # submodule would trip DefaultModelLoader.track_weights_loading, + # and a model-wide dtype/device pass would change VAE numerics. + self.__dict__["_vae"] = YuE2VAE.from_pretrained(vae_path, decoder_only=True, device="cpu") + if self._vae_device != device: + self._vae.to(device) + self._vae_device = device + return self._vae + + def _decode_latents_iter(self, latents: torch.Tensor) -> Iterator[None]: + """Device latents -> device stereo [2, samples] as a work generator + (one unit per VAE tile); nothing is read back to the host.""" + model = self._vae_model(latents.device) + audio = yield from model.decode_tiled_iter( + latents.T.unsqueeze(0), + core_frames=VAE_CORE_FRAMES, + halo_frames=VAE_HALO_FRAMES, + output_device=latents.device, + check_finite=False, + ) + # _SynthesisJob checks finite values before clipping and transferring. + return audio[0].contiguous() + + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + """[frames, 64] FP32 CPU latents -> channels-first stereo [2, samples].""" + device = f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu" + model = self._vae_model(torch.device(device)) + z = latents.T.unsqueeze(0) # [1, 64, T] + with torch.inference_mode(): + audio = model.decode_tiled( + z, + core_frames=VAE_CORE_FRAMES, + halo_frames=VAE_HALO_FRAMES, + output_device="cpu", + ) + if not torch.isfinite(audio).all(): + raise RuntimeError("VAE produced non-finite audio") + logger.debug("YUE2_AUDIO_PRECLAMP absmax=%.4f", float(audio.abs().max())) + # Channels-first [2, T], the same layout MiniMax Music 3 ships in + # model_outputs; serving's create_audio transposes it for soundfile. + return audio[0].float().clamp(-1, 1).contiguous() + # ------------------------------------------------------------ weights def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -1398,12 +2138,15 @@ class Yue2ForCausalLM(nn.Module): loader = AutoWeightsLoader(self) loaded = loader.load_weights(iter(ar_pairs)) + self.nar_fused = fuse_nar_projections(self.nar_layers) # Load the VAE now, not lazily at the first finishing request: a bad # path or a failed download must surface at startup, and the decoder # weights must sit on the GPU before vLLM's memory profiling sizes # the KV cache. device = f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu" self._vae_model(torch.device(device)) + self._warm_up_synthesis() + self._warm_up_sampler() # The NAR/projection modules were populated by hand above; # DefaultModelLoader.track_weights_loading diffs every named # parameter against the returned set, so the side keys must be @@ -1423,19 +2166,28 @@ class Yue2ForCausalLM(nn.Module): num_scheduled_tokens: Any, input_ids: torch.Tensor, positions: torch.Tensor, + sampling_extra_args: list[dict[str, Any]] | None = None, + discard_mask: Sequence[bool] | None = None, **_: Any, ) -> tuple[torch.Tensor, torch.Tensor]: - """Bind rows to requests for the step; prompt tokens are captured in - ``forward``, where the state is created (see ``_capture_constants``). + """Bind rows to requests and create new request state for the step. + + This runs before the backbone on every step, including CUDA graph + replays, so request admission never depends on forward() running. Decode rows feed tokens the model itself sampled, so they never contribute to the prompt prefix.""" computed = [int(v) for v in num_computed_tokens] scheduled = [int(v) for v in num_scheduled_tokens] self._step_rows = [(str(rid), computed[i], scheduled[i]) for i, rid in enumerate(req_ids)] + # The runner's own verdict on which rows' samples it keeps: a partial + # prefill, including a recompute after preemption (prompt + output + # tokens replayed), must not draw from the request's generator. + self._step_discard = None if discard_mask is None else [bool(v) for v in discard_mask] # make_omni_output runs every step; without this reset the first step's # payload would stay live forever and the _audio_queue fallback in # _ship_audio could never trigger. self._last_mm = None + self._capture_constants(sampling_extra_args, input_ids) return input_ids, positions def forward( @@ -1446,8 +2198,7 @@ class Yue2ForCausalLM(nn.Module): inputs_embeds: torch.Tensor | None = None, **kwargs: Any, ) -> torch.Tensor: - del intermediate_tensors - self._capture_constants(kwargs, input_ids) + del intermediate_tensors, kwargs hidden = self.model( input_ids=input_ids if inputs_embeds is None else None, positions=positions, @@ -1455,10 +2206,9 @@ class Yue2ForCausalLM(nn.Module): ) if isinstance(hidden, tuple): hidden = hidden[0] - self._flush_deferred_cleanup() return hidden - def _capture_constants(self, kwargs: dict[str, Any], input_ids: torch.Tensor | None) -> None: + def _capture_constants(self, extra_args: list[dict[str, Any]] | None, input_ids: torch.Tensor | None) -> None: """Create per-request state on the request's FIRST scheduled step. With KV prefix caching the first step may arrive with ``comp > 0`` @@ -1469,7 +2219,6 @@ class Yue2ForCausalLM(nn.Module): truth for the NAR conditioning prefix and for the prompt length the sampler uses to tell a completing prefill row from a mid-chunk one. """ - extra_args = kwargs.get("sampling_extra_args") if extra_args is None or input_ids is None: return offset = 0 @@ -1514,15 +2263,14 @@ class Yue2ForCausalLM(nn.Module): max_audio_frames=int(args.get(KEY_MAX_AUDIO_FRAMES, preset["max_tokens"])), seed=int(seed), skip_synthesis=bool(args.get(KEY_SKIP_SYNTHESIS, phase == "abc")), + max_hold_steps=int(args.get(KEY_MAX_HOLD_STEPS, 0) or 0), ) device = input_ids.device generator = torch.Generator(device=device if device.type != "mps" else "cpu") generator.manual_seed(constants.seed) - self._row_constants[req_id] = constants self._states[req_id] = _RequestState( request_id=req_id, constants=constants, - prompt_tokens=prompt_len, prompt_len=prompt_len, prefix_ids=prefix_ids, generator=generator, @@ -1530,65 +2278,349 @@ class Yue2ForCausalLM(nn.Module): offset += span def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor: - return self.logits_processor(self.lm_head, hidden_states) + if not self._semantic_only_step(): + return self.logits_processor(self.lm_head, hidden_states) + # Every row samples the semantic phase, whose sampler reads only the + # MUSIC_END + codec columns (18% of the vocabulary): project onto + # those alone and leave the rest of the row at -inf. + rows = hidden_states.shape[0] + if self._logits is None or self._logits.shape[0] < rows or self._logits.device != hidden_states.device: + self._logits = torch.full( + (max(rows, 8), self.vocab_size), -torch.inf, dtype=hidden_states.dtype, device=hidden_states.device + ) + logits = self._logits[:rows] + lo, hi = MUSIC_END, CODEC_OFFSET + CODEC_SIZE + torch.matmul(hidden_states, self.lm_head.weight[lo:hi].T, out=logits[:, lo:hi]) + return logits + + def _semantic_only_step(self) -> bool: + """True when every request in the step is a semantic-phase request.""" + if not self._step_rows or not hasattr(self.lm_head, "weight"): + return False + for req_id, _, _ in self._step_rows: + state = self._states.get(req_id) + if state is None or state.constants.phase != "semantic": + return False + return True # ------------------------------------------------------------ sampler def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> SamplerOutput | None: - """Draw one token per decoding row with the request's own arithmetic.""" + """Draw one token per decoding row with the request's own arithmetic. + + The draw stays on the device: its ids reach the host through an + async copy that the NEXT call resolves, so the host never waits for + the GPU inside a decode step and async scheduling can overlap the + next step's preparation with this step's kernels. Every draw still + sees the full history (the previous draw is resolved first); only + acting on an end token or a full frame budget moves one step later. + """ del sampling_metadata rows = int(logits.shape[0]) # The runner's input_ids buffer is int32; long ids crash its scatter # when concurrent requests re-index rows. - token_ids = torch.zeros((rows, 1), dtype=torch.int32, device=logits.device) + # Build fixed HOLD/end outputs on the host and transfer the whole batch + # once. Assigning a Python scalar to a CUDA row performs a blocking H2D + # copy. The pinned allocator tracks this copy's lifetime, so the buffer + # can be retired without waiting for it or rewriting an in-flight copy. + host_token_ids = torch.zeros((rows, 1), dtype=torch.int32, pin_memory=logits.is_cuda) + self._resolve_pending() + self._synthesis.pump() if not self._step_rows: - return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + return SamplerOutput( + sampled_token_ids=host_token_ids.to(logits.device, non_blocking=True), logprobs_tensors=None + ) + # Rows sharing a phase and preset are sampled as one batch. + groups: dict[tuple, list[tuple[int, _RequestState]]] = {} for row, (req_id, comp, span) in enumerate(self._step_rows): if row >= rows: break state = self._states.get(req_id) - if state is None or state.finished: + if state is None: continue - if comp + span < state.prompt_len: - # Mid-chunk prefill rows: their sampled token is discarded. - # (A completing prefill row reaches the full prompt length - # even when a prefix-cache hit left comp > 0.) + if self._step_discard is not None and row < len(self._step_discard): + discarded = self._step_discard[row] + else: + # A completing prefill row reaches the full prompt length even + # when a prefix-cache hit left comp > 0. + discarded = comp + span < state.prompt_len + if discarded: + continue + accepted = max(0, comp + span - state.prompt_len) + self._reconcile_history(state, accepted) + c = state.constants + end = ABC_END if c.phase == "abc" else MUSIC_END + if state.finished: + # Normally an async lookahead row the engine discards. If the + # scheduler dropped the step that emitted the end token (the + # request was preempted with it in flight), that step's audio + # was dropped too: deliver both again. + if state.end_len is not None and accepted < state.end_len: + if state.delivered is not None: + audio, truncated, error = state.delivered + self._ship_audio(state.request_id, audio, truncated, error=error) + state.end_len = accepted + 1 + host_token_ids[row, 0] = end + continue + if state.finish_ready: + state.finish_ready = False + if self._hold_allowed(state): + self._queue_synthesis(state) + host_token_ids[row, 0] = HOLD_TOKEN + continue + state.finished = True + state.end_len = accepted + 1 + self._finish_request_safely(state, hit_end=not state.truncated) + host_token_ids[row, 0] = end + continue + if state.job is not None: + if state.job.done() or not self._hold_allowed(state): + self._complete_synthesis(state) + state.end_len = accepted + 1 + host_token_ids[row, 0] = end + else: + state.hold_steps += 1 + host_token_ids[row, 0] = HOLD_TOKEN continue - constants = state.constants if state.request_id not in self._decode_t0: self._decode_t0[state.request_id] = time.perf_counter() - scores = distribution( - logits[row], - temperature=constants.temperature, - top_p=constants.top_p, - top_k=constants.top_k, - repetition_penalty=constants.repetition_penalty, - penalty_window=constants.penalty_window, - history=state.history, - step=len(state.history), - min_tokens=constants.min_tokens, - phase=constants.phase, + # With no penalty the window has no effect. Canonicalize it so + # equivalent groups cannot replay and overwrite the same graph's + # output buffers before their draws are collected below. + window = 0 if c.repetition_penalty == 1.0 else c.penalty_window + key = (c.phase, c.temperature, c.top_p, c.top_k, c.repetition_penalty, window) + groups.setdefault(key, []).append((row, state)) + token_ids = host_token_ids.to(logits.device, non_blocking=True) + if not groups: + return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + + # Stage every small host input of this step in one pinned buffer and + # move it with a single async copy: a pageable copy or a list index + # would make the host wait for the GPU. The previous step's copies + # have landed (_resolve_pending waited on its event), so the staging + # buffer can be rewritten. + plans = [] + flat: list[int] = [] + for (phase, temperature, top_p, top_k, penalty, window), members in groups.items(): + vocab = PHASE_VOCAB[phase] + # With no penalty there is no history input: [-0:] would walk the + # entire history every step before discarding it at width=0. + windows = ( + [] if penalty == 1.0 else [[vocab.col(t) for t in state.history[-window:]] for _, state in members] ) - token = sample_row(scores, state.generator, greedy=constants.temperature == 0) - end = ABC_END if constants.phase == "abc" else MUSIC_END - if token == end: - state.finished = True - state.truncated = False - self._finish_request_safely(state, hit_end=True) - token_ids[row, 0] = end + # A fixed width (the preset window) keeps the sampler graph static. + width = 0 if penalty == 1.0 else window if window > 0 else max(map(len, windows)) + block = [len(state.history) < state.constants.min_tokens for _, state in members] + # A request that still needs its song must not stop on its end + # token when drawn: the host only sees the draw next step. + end = ABC_END if phase == "abc" else MUSIC_END + held = [-1 if state.constants.skip_synthesis else end for _, state in members] + base = len(flat) + flat += [row for row, _ in members] + block + held + for cols in windows if width else (): + flat += cols + [vocab.num_cols] * (width - len(cols)) + plans.append((vocab, (temperature, top_p, top_k, penalty), members, base, width, any(block))) + if self._stage is None or self._stage.numel() < len(flat): + self._stage = torch.empty(max(len(flat), 1024), dtype=torch.long, pin_memory=True) + self._stage[: len(flat)].copy_(torch.as_tensor(flat)) + staged = self._stage[: len(flat)].to(logits.device, non_blocking=True) + + # Graphs share a pool and may replay in a different order from capture. + # Preserve each group's outputs outside that pool before replaying the + # next graph: its intermediates may alias the previous graph's outputs. + count = sum(len(members) for members in groups.values()) + drawn = torch.empty((2, count), dtype=torch.long, device=logits.device) + emitted: list[torch.Tensor] = [] + drawn_rows: list[int] = [] + states: list[_RequestState] = [] + for vocab, (temperature, top_p, top_k, penalty), members, base, width, blocks in plans: + n = len(members) + index = staged[base : base + n] + rows_logits = logits if n == rows else logits.index_select(0, index) + window_cols = staged[base + 3 * n : base + 3 * n + n * width].view(n, width) if width else None + block_end = staged[base + n : base + 2 * n].bool() if blocks else None + generators = [state.generator for _, state in members] + graph = self._sampler_graph(vocab, n, rows_logits, (temperature, top_p, top_k, penalty), width) + if graph is not None: + ids, bad = graph(rows_logits, window_cols, block_end, generators) else: - state.history.append(token) - budget = constants.max_audio_frames - if constants.phase == "semantic" and len(state.history) >= budget: - state.finished = True - state.truncated = True - self._finish_request_safely(state, hit_end=False) - token_ids[row, 0] = end - else: - token_ids[row, 0] = token + ids, bad = sample_rows( + rows_logits, + vocab, + temperature=temperature, + top_p=top_p, + top_k=top_k, + repetition_penalty=penalty, + window_cols=window_cols, + block_end=block_end, + generators=generators, + ) + offset = len(drawn_rows) + drawn[0, offset : offset + n].copy_(ids) + drawn[1, offset : offset + n].copy_(bad) + emitted.append(torch.where(ids == staged[base + 2 * n : base + 3 * n], HOLD_TOKEN, ids)) + drawn_rows += [row for row, _ in members] + states += [state for _, state in members] + out = torch.cat(emitted).to(torch.int32) + if drawn_rows == list(range(rows)): + token_ids[:, 0] = out + else: + index = torch.cat([staged[base : base + len(m)] for _, _, m, base, _, _ in plans]) + token_ids.view(-1).index_copy_(0, index, out) + + if self._pinned is None or self._pinned.shape[1] < count: + self._pinned = torch.empty((2, max(count, 64)), dtype=torch.long, pin_memory=True) + host = self._pinned[:, :count] + # Async into pinned memory; the next sample() call resolves it. + host.copy_(drawn, non_blocking=True) + event = torch.Event() + event.record() + self._pending = (event, host, states) return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + def _sampler_graph(self, vocab, rows, logits, preset, width) -> SamplerGraph | None: + """Use a power-of-two bucket; bound captures for custom sampling presets. + + The serving presets are captured before vLLM profiles memory. Other + presets run eagerly once the small custom-preset allowance is full. + """ + if logits.device.type != "cuda": + return None + bucket = 1 << (rows - 1).bit_length() + key = (vocab.phase, bucket, logits.dtype, preset, width) + graph = self._sampler_graphs.get(key) + if graph is None: + if len(self._sampler_graphs) >= 2 * self._sampler_capacity.bit_length() + 8: + return None + if self._sampler_graph_pool is None: + self._sampler_graph_pool = torch.cuda.graph_pool_handle() + graph = SamplerGraph( + vocab, bucket, logits.shape[1], logits.dtype, preset, width, logits.device, self._sampler_graph_pool + ) + self._sampler_graphs[key] = graph + return graph + + def _warm_up_sampler(self) -> None: + """Capture default phase/preset buckets before serving and KV profiling. + + Lazy captures use torch.cuda.graph, which synchronizes the device. + Admission and batch-size changes should replay pre-existing graphs. + """ + weight = self.lm_head.weight + if weight.device.type != "cuda": + return + bucket = 1 + while bucket < self._sampler_capacity * 2: + for phase, preset in (("semantic", SEMANTIC_SAMPLING), ("abc", ABC_SAMPLING)): + shape = (preset["temperature"], preset["top_p"], preset["top_k"], preset["repetition_penalty"]) + logits = torch.empty((bucket, self.vocab_size), dtype=weight.dtype, device=weight.device) + self._sampler_graph(PHASE_VOCAB[phase], bucket, logits, shape, preset["penalty_window"]) + if bucket >= self._sampler_capacity: + break + bucket *= 2 + + # ------------------------------------------------------------ async synthesis + + def _hold_allowed(self, state: _RequestState) -> bool: + """May this request emit one more HOLD_TOKEN step? + + Needs the caller's hold budget (and its max_tokens headroom) and room + in the context for the HOLD and the final end token. + """ + c = state.constants + if c.skip_synthesis or state.hold_steps >= c.max_hold_steps or not torch.cuda.is_available(): + return False + # prompt + frames + the drawn-end HOLD + holds so far + this HOLD + end + return state.prompt_len + len(state.history) + state.hold_steps + 3 <= self._max_model_len + + def _song_iter( + self, prefix_ids: list[int], codec: list[int], seed: int, keep: list + ) -> Generator[None, None, torch.Tensor]: + """NAR + VAE for one song as work units; returns the device audio.""" + latents = yield from synthesize_iter(self, prefix_ids, codec, seed, keep, self._synthesis.graph_pool) + keep.append(latents) + return (yield from self._decode_latents_iter(latents)) + + def _queue_synthesis(self, state: _RequestState) -> None: + codec = [t - CODEC_OFFSET for t in state.history] + prefix_ids, seed = state.prefix_ids, state.constants.seed + job = _SynthesisJob(state.request_id, lambda keep: self._song_iter(prefix_ids, codec, seed, keep)) + if not codec or min(codec) < 0 or max(codec) >= CODEC_SIZE: + job.error = RuntimeError("semantic history is empty or contains non-codec tokens") + state.job = job + state.hold_steps = 1 + self._synthesis.submit(job) + + def _complete_synthesis(self, state: _RequestState) -> None: + """Ship the song on this step (waiting for it if the hold ran out).""" + job = state.job + assert job is not None + audio, failed, waited = self._synthesis.complete(job) + state.job = None + state.finished = True + t0 = self._decode_t0.pop(state.request_id, None) + t_ar = (job.t_start - t0) if t0 is not None and job.t_start is not None else float("nan") + logger.debug( + "YUE2_TIMING req=%s phase=%s ar_tokens=%d ar_s=%.2f ar_tps=%.1f synth_s=%.2f hold_steps=%d waited=%s", + state.request_id, + state.constants.phase, + len(state.history), + t_ar, + len(state.history) / t_ar if t_ar else float("nan"), + (job.t_done - job.t_start) if job.t_start is not None and job.t_done is not None else 0.0, + state.hold_steps, + waited, + ) + self._ship_audio(state.request_id, audio, state.truncated, error=failed) + + def _reconcile_history(self, state: _RequestState, accepted: int) -> None: + """Drop in-flight draws the scheduler discarded on preemption. + + Only a non-discarded row has replayed the complete accepted sequence; + partial recompute chunks cannot establish this boundary. A drawn EOS + occupies one engine token even though it is not in the codec history. + """ + if accepted >= len(state.history) + int(state.end_drawn): + return + if state.job is not None: + self._synthesis.cancel(state.job) + state.job = None + del state.history[accepted:] + state.end_drawn = False + state.finished = False + state.end_len = None + state.delivered = None + state.hold_steps = 0 + state.truncated = state.constants.phase == "semantic" and len(state.history) >= state.constants.max_audio_frames + state.finish_ready = state.truncated + + def _resolve_pending(self) -> None: + """Apply the previous step's draws to the host-side request state.""" + if self._pending is None: + return + event, host, states = self._pending + self._pending = None + event.synchronize() + ids, bad = host.tolist() + if any(bad): + raise RuntimeError("sampling scores are all -inf; check phase masking") + for state, token in zip(states, ids): + constants = state.constants + if token == (ABC_END if constants.phase == "abc" else MUSIC_END): + state.end_drawn = True + state.truncated = False + if constants.skip_synthesis: + # Emitted as drawn: the engine has already stopped it. + state.finished = True + else: + state.finish_ready = True + continue + state.history.append(token) + if constants.phase == "semantic" and len(state.history) >= constants.max_audio_frames: + state.truncated = True + state.finish_ready = True + # ------------------------------------------------------------ audio def _finish_request_safely(self, state: _RequestState, *, hit_end: bool) -> None: @@ -1602,7 +2634,9 @@ class Yue2ForCausalLM(nn.Module): """ try: self._finish_request(state, hit_end=hit_end) - except Exception: + except (RuntimeError, ValueError, FloatingPointError): + # CPU synthesis also reports non-finite latents as FloatingPointError. + # Programming errors propagate on both CUDA completion paths. logger.exception("YuE2 synthesis failed for req=%s; failing only this request", state.request_id) self._ship_audio(state.request_id, torch.zeros((2, 0)), state.truncated, error=True) @@ -1611,25 +2645,22 @@ class Yue2ForCausalLM(nn.Module): constants = state.constants if constants.skip_synthesis or not state.history: return + if torch.cuda.is_available(): + # No hold headroom means wait now, but still use the serialized + # synthesis stream and its bounded in-flight work. Independent sync + # graphs would overlap a running job and accumulate private pools. + self._queue_synthesis(state) + state.hold_steps = 0 + self._complete_synthesis(state) + return codec = [t - CODEC_OFFSET for t in state.history] if min(codec) < 0 or max(codec) >= CODEC_SIZE: raise RuntimeError("semantic history contains non-codec tokens") t0 = self._decode_t0.pop(state.request_id, None) t_nar0 = time.perf_counter() - latents = synthesize(self, state.prefix_ids, codec, constants.seed, steps=ODE_STEPS) - logger.info( - "YUE2_LATENTS frames=%d absmax=%.4f mean=%.5f std=%.4f", - latents.shape[0], - float(latents.abs().max()), - float(latents.mean()), - float(latents.std()), - ) - t_nar = time.perf_counter() - t_nar0 - t_vae0 = time.perf_counter() - audio = self._decode_latents(latents) - t_vae = time.perf_counter() - t_vae0 + audio, t_nar, t_vae = self.synthesize_song(state.prefix_ids, codec, constants.seed) t_ar = (t_nar0 - t0) if t0 is not None else float("nan") - logger.info( + logger.debug( "YUE2_TIMING req=%s phase=%s ar_tokens=%d ar_s=%.2f ar_tps=%.1f nar_s=%.2f vae_s=%.2f", state.request_id, constants.phase, @@ -1651,7 +2682,12 @@ class Yue2ForCausalLM(nn.Module): sparse routing on the request's last step in the output batch. Leaving the payload for a later step would drop it: the request is gone by then (gepard_talker flushes at its last step for the same reason). + The request keeps a reference until the engine finishes it, in case + the scheduler drops this step (see ``_RequestState.end_len``). """ + state = self._states.get(req_id) + if state is not None: + state.delivered = (audio, truncated, error) mm = self._last_mm if mm is None: # No forward built a payload this step; fall back to the queue. @@ -1666,39 +2702,6 @@ class Yue2ForCausalLM(nn.Module): meta["truncated"].append(int(truncated)) meta["error"].append(int(error)) - def _vae_model(self, device: torch.device) -> YuE2VAE: - if self._vae is None: - vae_path = os.environ.get("YUE2_VAE", DEFAULT_VAE_ID) - logger.info("Loading YuE2 VAE decoder from %s", vae_path) - # Keep the VAE out of the module tree (__dict__, not an nn.Module - # attribute): the checkpoint has no VAE keys, so a registered - # submodule would trip DefaultModelLoader.track_weights_loading, - # and a model-wide dtype/device pass would change VAE numerics. - self.__dict__["_vae"] = YuE2VAE.from_pretrained(vae_path, decoder_only=True, device="cpu") - if self._vae_device != device: - self._vae.to(device) - self._vae_device = device - return self._vae - - def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: - """[frames, 64] FP32 CPU latents -> channels-first stereo [2, samples].""" - device = f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu" - model = self._vae_model(torch.device(device)) - z = latents.T.unsqueeze(0) # [1, 64, T] - with torch.inference_mode(): - audio = model.decode_tiled( - z, - core_frames=VAE_CORE_FRAMES, - halo_frames=VAE_HALO_FRAMES, - output_device="cpu", - ) - if not torch.isfinite(audio).all(): - raise RuntimeError("VAE produced non-finite audio") - logger.info("YUE2_AUDIO_PRECLAMP absmax=%.4f", float(audio.abs().max())) - # Channels-first [2, T], the same layout MiniMax Music 3 ships in - # model_outputs; serving's create_audio transposes it for soundfile. - return audio[0].float().clamp(-1, 1).contiguous() - def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: if isinstance(model_outputs, OmniOutput): return model_outputs @@ -1733,18 +2736,20 @@ class Yue2ForCausalLM(nn.Module): return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: - # Fires before forward; defer the free so the in-flight step can read. - for rid in finished_req_ids: - self._deferred_cleanup_ids.add(_request_key(rid)) - - def _flush_deferred_cleanup(self) -> None: - for req_id in self._deferred_cleanup_ids: + # Runs at the start of execute_model, after the previous step's + # sample(), so the state can be freed right away. + finished = {_request_key(rid) for rid in finished_req_ids} + for req_id in finished: state = self._states.pop(req_id, None) + if state is not None and state.job is not None: + # Aborted mid-synthesis. + self._synthesis.cancel(state.job) + state.job = None if state is not None and not state.finished: logger.warning( "YuE2 request %s ended without an end token or its frame " "budget (aborted or preempted); no audio was produced.", req_id, ) - self._row_constants.pop(req_id, None) - self._deferred_cleanup_ids.clear() + self._decode_t0.pop(req_id, None) + self._step_rows = [row for row in self._step_rows if row[0] not in finished] diff --git a/vllm_omni/tokenizers/yue2_prompt.py b/vllm_omni/tokenizers/yue2_prompt.py index 97d7b7dea..f1baa4122 100644 --- a/vllm_omni/tokenizers/yue2_prompt.py +++ b/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] diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 330355c64..63dde07c2 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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],